Sparse Layers

class weavenet.sparse.layers.MaskSelectorByLinearInferenceOr(dim_src=-3, dim_tar=-2, drop_rate=0.5, tau=1.0)

Selects edges based on linear prediction. The result for each direction is aggregated by OR rule.

Parameters

• dim_src (int) – dim of source vertex of edges.

• dim_tar (int) – dim of target vertex of edges.

• drop_rate (float) – sets drop rate of edges for each vertex. The OR rule selection may results in less drop-rate in actual calculations.

• tau (float) – the temperature of gumbel sigmoid.

build(input_channels, output_channels=1)

Build the linear layer for the prediction. This function is automatically called in TrainableMatchingModuleSp.

Parameters

• input_channels (int) – the number of input channels.

• output_channels (int) – the number of output channels.

Return type

None

forward(xab, xba_t)

Shape:

• xab: \((B, N, M, C)\)

• xba_t: \((B, N, M, C)\)

• output: \((B, N, M, C')\), where \(C'\) is typically 1.

Parameters

• xab (Tensor) – batched feature map, typically with the size of (B, N, M, C) where ij-th feature at \((i, j)\in N \times M\) represent edges from side a to b.

• xba_t (Tensor) – batched feature map with the same shape with xab, and represent edges from side b to a.

Returns

• mask, where edges with score 1.0 are selected and 0.0 are dropped.

Return type

Tensor

class weavenet.sparse.layers.MaskSelectorByNorm(dim_src=-3, dim_tar=-2, drop_rate=0.5, tau=1.0)

Selects edges based on linear prediction. The result for each direction is aggregated by OR rule.

Parameters

• dim_src (int) – dim of source vertex of edges.

• dim_tar (int) – dim of target vertex of edges.

• drop_rate (float) – sets drop rate of edges for each vertex. The OR rule selection may results in less drop-rate in actual calculations.

• tau (float) – the temperature of gumbel sigmoid.

forward(xab, xba_t)

Shape:

• xab: \((B, N, M, C)\)

• xba_t: \((B, N, M, C)\)

• output: \((B, N, M, C')\), where \(C'\) is typically 1.

Parameters

• xab (Tensor) – batched feature map, typically with the size of (B, N, M, C) where ij-th feature at \((i, j)\in N \times M\) represent edges from side a to b.

• xba_t (Tensor) – batched feature map with the same shape with xab, and represent edges from side b to a.

Returns

• mask, where edges with score 1.0 are selected and 0.0 are dropped.

Return type

Tensor

class weavenet.sparse.layers.MaskSelectorBySimilarity(max_edge_survive_rate=1.1, max_survive_edges_per_sample=-1)

select mask for edge pruning with radius neighbor.

Parameters

• radius – the threshold of radius neighbor. If set a value <=0, only the max_edge_survive_rate is considered. This <=0 option may accelerate the process by fixing the network shape [default: 0].

• max_edge_survive_rate (float) – set maximum number of edges selected in this selector by rate [default: 1.0]. e.g., a 4x4 graph has 16 edges, and at most 7 will path if the rate is 0.5.

• max_suvive_edges_per_sample – set the maximum number of edges selected in this selector by the number of edges per sample. If <=0, ignored. If <N (or <M), set N (or M)[default: -1]

• max_survive_edges_per_sample (int) –

get_threshold_by_k(sim, dim=-1, k=-1)

calculate threshold that prevend memory overflow, based on the params.

Shape:

• sim: (ldots, N, M)

Parameters

• sim (Tensor) – a similarity matrix.

• dim (int) – dimention of kth value.

• k (int) –

Returns

the number of edges survive.

Return type

Tensor

wrapup(sim, sim_selected)

Parameters

• sim (Tensor) –

• sim_selected (Tensor) –

Return type

Tensor

class weavenet.sparse.layers.MaskSelectorRadiusNeighbor(radius=0.0, max_edge_survive_rate=1.1, max_survive_edges_per_sample=-1)

select mask for edge pruning with radius neighbor.

Parameters

• radius (float) – the threshold of radius neighbor. If set a value <=0, only the max_edge_survive_rate is considered. This <=0 option may accelerate the process by fixing the network shape [default: 0].

• max_edge_survive_rate (float) – see MaskSelectorBySimilarity

• max_suvive_edges_per_sample – see MaskSelectorBySimilarity

• max_survive_edges_per_sample (int) –

forward()

Shape:

• sim: (ldots, N, M)

Parameters

sim (Tensor) – a similarity matrix.

Returns

a differential mask. Elements with mask==1 is selected.

Return type

Tensor

class weavenet.sparse.layers.MaskSelectorReciprocalNeighbor(k=0, max_edge_survive_rate=0.5, max_suvive_edges_per_sample=-1)

Select mask for edge pruning with \(k\)-reciprocal neighbors. Aｎ element \(x_a\) (on side a) is k-reciprocal neighbor of \(x_b\) (on side b) when \(x_a\) is in top-\(k\) neighbors of \(x_b\) and vice .

Parameters

• k (int) – use k-th nearest vertices as neighbors. If <=0, ignored [default: 0]

• max_edge_survive_rate (float) – see MaskSelectorBySimilarity

• max_suvive_edges_per_sample (int) – see MaskSelectorBySimilarity

forward()

Shape:

• sim: (ldots, N, M)

Parameters

sim (Tensor) – a similarity matrix.

Returns

a differential mask. Elements with mask==1 is selected.

Return type

Tensor

class weavenet.sparse.layers.SimilarityBasedMaskInference(compute_similarity=<torch.jit.ScriptFunction object>, mask_selector=MaskSelectorRadiusNeighbor())

Inferences mask based on similarity.

Parameters

• compute_similarity (Callable[[Tensor, Tensor], Tensor]) – a callable object that calculates a similarity matrix.

• select_mask – a callable object that select a (differentiable) mask based on the similarity. [Default: :class:`MaskSelectorRadiusNeighbor`(0.0,0.5)].

• mask_selector (Callable[[Tensor], Tensor]) –

forward(xab, xba_t)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters

• xab (Tensor) –

• xba_t (Tensor) –

Return type

Tensor

class weavenet.sparse.layers.SparseDenseAdaptor(mask)

Adapt sparse-dense matrix conversion, based on a given mask.

Shape:

• mask: (ldots, N, M, 1)

Parameters

mask (Tensor) – a mask that selects edges.

to_sparse(x)

Shape:

• x: (ldots, N, M, C)

• output: \((\text{num_of_selected_edges_in_batch}, C)\)

Parameters

x (Tensor) – batched edge features.

Returns

a flatten edge features, whose elements are selected by mask.

Return type

Tensor

to_dense(x_sparse)

Shape:

• x_sparse: \((\text{num_of_selected_edges_in_batch}, C)\)

• output: \((\dots, N, M, C)\)

Parameters

x_sparse (Tensor) – a flattend edge features.

Returns

a edge features reformatted in the original shape.

Return type

Tensor

class weavenet.sparse.layers.MaxPoolingAggregatorSp

forward(x_sp, vertex_id, dim=0)

Shape:

• x: \((\ldots, M, D)\) if dim = -2, otherwise, the axis directed by dim should have M and aggregated while keeping dims.

• output: \((\ldots, 1, D)\)

Parameters

• x – an input tensor.

• x_sp (Tensor) –

• vertex_id (Tensor) –

• dim (int) –

Returns

x_aggregated

Return type

Tensor

class weavenet.sparse.layers.SetEncoderBaseSp(first_process, aggregator, second_process_edge, second_process_vertex)

A sparse version of SetEncoderBase

Parameters

• first_process (Callable[[Tensor], Tensor]) – a callable (and typically trainable) object that converts a \((B, N, M, C_{input})\) tensor to \((B, N, M, C_{mid})\).

• aggregator (Callable[[Tensor, Tensor], Tensor]) – a callable object that aggregate \(M\) edge features for each of \(N\) vertices. The resultant tensor is reformatted into the shape of \((B, N, M, C_{mid})\) tensor.

• merger – a callable object that merge \((B, N, M, C_{input})\) edge features and \((B, N, M, C_{mid})\) vertex features into \((B, N, M, C_{merged})\).

• second_process – a callable (and typically trainable) object that converts a \((B, N, M, C_{merged})\) tensor to \((B, N, M, C_{output})\).

• second_process_edge (Callable[[Tensor], Tensor]) –

• second_process_vertex (Callable[[Tensor], Tensor]) –

forward(x, vertex_id)

Applies set encoding operations.

Shape:

• x: \((\text{num_of_edges_in_batch}, \text{in_channels})\)

• vertex_id: \((\text{num_of_edges_in_batch}, )\)

• output: \((\text{num_of_edges_in_batch}, \text{output_channels})\)

Parameters

• x (Tensor) – an input tensor.

• vertex_id (Tensor) – an index list of vertex id for each edge.

Returns

x_processed

Return type

Tensor

class weavenet.sparse.layers.SetEncoderPointNetSp(in_channels, mid_channels, output_channels, **kwargs)

A sparse version of SetEncoderPointNet

Parameters

• in_channels (int) – the number of input channels.

• mid_channels (int) – the number of output channels at the first convolution.

• out_channels – the number of output channels at the second convolution.

• output_channels (int) –

class weavenet.sparse.layers.DualSoftmaxSp

A sparse version of DualSoftmax

apply_softmax(xab, src_id, tar_id, xba=None)

Parameters

• xab (Tensor) –

• src_id (Tensor) –

• tar_id (Tensor) –

• xba (Optional[Tensor]) –

Return type

Tuple[Tensor, Tensor]

forward(xab, src_id, tar_id, xba=None)

Calculate the dual softmax for batched matrices.

Shape:

• xab: \((\text{num_of_edges_in_batch}, \text{in_channels})\)

• src_id: \((\text{num_of_edges_in_batch}, )\)

• tar_id: \((\text{num_of_edges_in_batch}, )\)

• xba: \((\text{num_of_edges_in_batch}, \text{in_channels})\)

• output: \((\text{num_of_edges_in_batch}, \text{in_channels})\) (all the three outputs has the same shape).

Parameters

• xab (Tensor) – 1st batched matrices.

• src_id (Tensor) – an index list of source vertex id for each edge.

• tar_id (Tensor) – an index list of target vertex id for each edge.

• xba (Optional[Tensor]) – 2nd batched matrices. If None, xab is used as xba.

Returns

a triplet of (mab * mba), mab (=softmax(xab, dim=-2)), mba (=softmax(xba_t, dim=-1)

Return type

Tuple[Tensor, Tensor, Tensor]

class weavenet.sparse.layers.DualSoftmaxSqrtSp

A sparse version of DualSoftmaxSqrt

forward(xab, src_id, tar_id, xba=None)

Shape and Args: same as DualSoftmaxSp

Args:

xab: 1st batched matrices. src_id: an index list of source vertex id for each edge. tar_id: an index list of target vertex id for each edge. xba: 2nd batched matrices. If None, xab is used as (transposed) xba. This option corresponds to the original implementation of LoFTR’s dual softmax.

Returns

values (mab * mba_t).sqrt(), mab (=softmax(xab, dim=-2)), mba_t (=softmax(xba_t, dim=-1)

Parameters

• xab (Tensor) –

• src_id (Tensor) –

• tar_id (Tensor) –

• xba (Optional[Tensor]) –

Return type

Tuple[Tensor, Tensor, Tensor]

class weavenet.sparse.layers.DualSoftmaxFuzzyLogicAndSp

A sparse version of DualSoftmaxFuzzyLogicAnd

forward(xab, xba=None, is_xba_transposed=True)

Shape and Args: same as DualSoftmaxSp

Args:

xab: 1st batched matrices. src_id: an index list of source vertex id for each edge. tar_id: an index list of target vertex id for each edge. xba: 2nd batched matrices. If None, xab is used as (transposed) xba. This option corresponds to the original implementation of LoFTR’s dual softmax.

Returns

values torch.min(mab, mba_t), mab (=softmax(xab, dim=-2)), mba_t (=softmax(xba_t, dim=-1)

Parameters

• xab (Tensor) –

• xba (Optional[Tensor]) –

• is_xba_transposed (bool) –

Return type

Tuple[Tensor, Tensor, Tensor]

class weavenet.sparse.layers.SparseDenseAdaptor(mask)

Adapt sparse-dense matrix conversion, based on a given mask.

Shape:

• mask: (ldots, N, M, 1)

Parameters

mask (Tensor) – a mask that selects edges.

to_sparse(x)

Shape:

• x: (ldots, N, M, C)

• output: \((\text{num_of_selected_edges_in_batch}, C)\)

Parameters

x (Tensor) – batched edge features.

Returns

a flatten edge features, whose elements are selected by mask.

Return type

Tensor

to_dense(x_sparse)

Shape:

• x_sparse: \((\text{num_of_selected_edges_in_batch}, C)\)

• output: \((\dots, N, M, C)\)

Parameters

x_sparse (Tensor) – a flattend edge features.

Returns

a edge features reformatted in the original shape.

Return type

Tensor

class weavenet.sparse.layers.MaxPoolingAggregatorSp

forward(x_sp, vertex_id, dim=0)

Shape:

• x: \((\ldots, M, D)\) if dim = -2, otherwise, the axis directed by dim should have M and aggregated while keeping dims.

• output: \((\ldots, 1, D)\)

Parameters

• x – an input tensor.

• x_sp (Tensor) –

• vertex_id (Tensor) –

• dim (int) –

Returns

x_aggregated

Return type

Tensor

class weavenet.sparse.layers.SetEncoderBaseSp(first_process, aggregator, second_process_edge, second_process_vertex)

A sparse version of SetEncoderBase

Parameters

• first_process (Callable[[Tensor], Tensor]) – a callable (and typically trainable) object that converts a \((B, N, M, C_{input})\) tensor to \((B, N, M, C_{mid})\).

• aggregator (Callable[[Tensor, Tensor], Tensor]) – a callable object that aggregate \(M\) edge features for each of \(N\) vertices. The resultant tensor is reformatted into the shape of \((B, N, M, C_{mid})\) tensor.

• merger – a callable object that merge \((B, N, M, C_{input})\) edge features and \((B, N, M, C_{mid})\) vertex features into \((B, N, M, C_{merged})\).

• second_process – a callable (and typically trainable) object that converts a \((B, N, M, C_{merged})\) tensor to \((B, N, M, C_{output})\).

• second_process_edge (Callable[[Tensor], Tensor]) –

• second_process_vertex (Callable[[Tensor], Tensor]) –

forward(x, vertex_id)

Applies set encoding operations.

Shape:

• x: \((\text{num_of_edges_in_batch}, \text{in_channels})\)

• vertex_id: \((\text{num_of_edges_in_batch}, )\)

• output: \((\text{num_of_edges_in_batch}, \text{output_channels})\)

Parameters

• x (Tensor) – an input tensor.

• vertex_id (Tensor) – an index list of vertex id for each edge.

Returns

x_processed

Return type

Tensor

class weavenet.sparse.layers.SetEncoderPointNetSp(in_channels, mid_channels, output_channels, **kwargs)

A sparse version of SetEncoderPointNet

Parameters

• in_channels (int) – the number of input channels.

• mid_channels (int) – the number of output channels at the first convolution.

• out_channels – the number of output channels at the second convolution.

• output_channels (int) –

class weavenet.sparse.layers.DualSoftmaxSp

A sparse version of DualSoftmax

apply_softmax(xab, src_id, tar_id, xba=None)

Parameters

• xab (Tensor) –

• src_id (Tensor) –

• tar_id (Tensor) –

• xba (Optional[Tensor]) –

Return type

Tuple[Tensor, Tensor]

forward(xab, src_id, tar_id, xba=None)

Calculate the dual softmax for batched matrices.

Shape:

• xab: \((\text{num_of_edges_in_batch}, \text{in_channels})\)

• src_id: \((\text{num_of_edges_in_batch}, )\)

• tar_id: \((\text{num_of_edges_in_batch}, )\)

• xba: \((\text{num_of_edges_in_batch}, \text{in_channels})\)

• output: \((\text{num_of_edges_in_batch}, \text{in_channels})\) (all the three outputs has the same shape).

Parameters

• xab (Tensor) – 1st batched matrices.

• src_id (Tensor) – an index list of source vertex id for each edge.

• tar_id (Tensor) – an index list of target vertex id for each edge.

• xba (Optional[Tensor]) – 2nd batched matrices. If None, xab is used as xba.

Returns

a triplet of (mab * mba), mab (=softmax(xab, dim=-2)), mba (=softmax(xba_t, dim=-1)

Return type

Tuple[Tensor, Tensor, Tensor]

class weavenet.sparse.layers.DualSoftmaxSqrtSp

A sparse version of DualSoftmaxSqrt

forward(xab, src_id, tar_id, xba=None)

Shape and Args: same as DualSoftmaxSp

Args:

xab: 1st batched matrices. src_id: an index list of source vertex id for each edge. tar_id: an index list of target vertex id for each edge. xba: 2nd batched matrices. If None, xab is used as (transposed) xba. This option corresponds to the original implementation of LoFTR’s dual softmax.

Returns

values (mab * mba_t).sqrt(), mab (=softmax(xab, dim=-2)), mba_t (=softmax(xba_t, dim=-1)

Parameters

• xab (Tensor) –

• src_id (Tensor) –

• tar_id (Tensor) –

• xba (Optional[Tensor]) –

Return type

Tuple[Tensor, Tensor, Tensor]

Experimental

class weavenet.sparse.layers.DualSoftmaxFuzzyLogicAndSp

A sparse version of DualSoftmaxFuzzyLogicAnd

forward(xab, xba=None, is_xba_transposed=True)

Shape and Args: same as DualSoftmaxSp

Args:

xab: 1st batched matrices. src_id: an index list of source vertex id for each edge. tar_id: an index list of target vertex id for each edge. xba: 2nd batched matrices. If None, xab is used as (transposed) xba. This option corresponds to the original implementation of LoFTR’s dual softmax.

Returns

values torch.min(mab, mba_t), mab (=softmax(xab, dim=-2)), mba_t (=softmax(xba_t, dim=-1)

Parameters

• xab (Tensor) –

• xba (Optional[Tensor]) –

• is_xba_transposed (bool) –

Return type

Tuple[Tensor, Tensor, Tensor]

Private Layers