Layers

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

Applies a process proposed in DeepSet and PointNet to a set of out-edge features of each vertex. See SetEncoderBase for its forwarding process.

Parameters

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

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

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

class weavenet.layers.CrossConcat(dim_feature=-1)

Applies cross-concatenation introduced in Edge-Selective Feature Weaving for Point Cloud Matching

\[\text{CrossConcat}([x^{ab}, {x^{ba}}^\top]) = [\text{cat}([x^{ab}, {x^{ba}}^\top], dim=-1), \text{cat}([{x^{ba}}^\top,x^{ab}], dim=-1)]\]

Example of Usage:

# prepare an instance of this class.
interactor = CrossConcat()

# assume batch_size=8, the problem instance size is 6x5,  and each edge feature is 32 channels.
B, N, M, C = 8, 6, 5, 32

xab = torch.rand((B, N, M, C), dtype=torch.float) # prepare a random input. # NxM
xba = torch.rand((B, M, N, C), dtype=torch.float) # prepare a random input. # MxN
xba_t = xba.transpose(1,2)
zab, zba_t = interactor(xab, xba_t)
assert(xab.size(-1)*2 == 2*C)
assert(xba_t.size(-1)*2 == 2*C)
assert(interactor.output_channels(C)==2*C)

Parameters

dim_feature (int) –

output_channels(input_channels)

Parameters

input_channels (int) – assumed input channels.

Returns

\(2*\) input_channels

Return type

output_channels

forward(xab, xba_t)

Shape:

• xab: \((\ldots, C)\)

• xba_t: \((\ldots, C)\)

• output(zab, zba_t): \([(\ldots, 2*C),(\ldots, 2*C)]\)

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

(zab, zba_t) calculated as \(\text{CrossConcat}([x^{ab}, {x^{ba}}^\top])\) .

Return type

Tuple[Tensor, Tensor]

class weavenet.layers.CrossConcatVertexFeatures(dim_a=-3, dim_b=-2, dim_feature=-1, compute_similarity=None, directional_normalization=None)

CrossConcat vertex features for side a and b, which are typically provided from a feature extractor.

Parameters

• dim_a – the dimension of side a in the output tensor.

• dim_b – the dimension of side b in the output tensor.

• compute_similarity (Optional[Callable[[Tensor, Tensor], Tensor]]) – an operator to calculate similarity between two vertex features. If set, the output has an additional channels where the similarity is set.

• directional_normalization (Optional[Callable[[Tensor, int], Tensor]]) – an normalizer applied to the axe dim_a and dim_b. If set, the output has two additional channels where normalized similarities on the two axe are set. If compute_similarity is None, this function will be ignored.

forward(xa, xb)

Concat vertex features on the two sides a and b, while appending similarities based on the callback functions.

Shape:

• x_a: \((\ldots, N, C)\)

• x_b: \((\ldots, M, C)\)

• output: \((\ldots, N, M, C')\) if dim_a = -3 and dim_b = -2. \(C' = 2*C\) if compute_similarity is None, \(C' = 2*C+1\) if only compute_similarity is set, and :math:`C’ = 2*C+2`if both compute_similarity and directional_normalization are set.

Parameters

• xa (Tensor) – vertex features on the side a.

• xb (Tensor) – vertex features on the side b.

Returns

A pair of edge-wise feature block tensors. Its feature consists of \(2*C\)-dimensional features (whose \(ij\)-th feature is a concatenation of :math:` and 1 or 2 dimensional similarities (depending on the callback function setting at __init__().

Return type

Tensor

output_channels(input_channels)

Parameters

input_channels (int) – assumed input channels.

Returns

\(2*\) input_channels if self.compute_similarity is None, \(2*\) input_channels \(+1\) if only self.compute_similarity is set, and \(2*\) input_channels :math:`+2`if both self.compute_similarity and self.directional_normalization are set.

Return type

output_channels

class weavenet.layers.MaxPoolingAggregator

Aggregates edge features for each vertex, and virtually reshape it to have the same size (other chann \(C\)) for merger.

forward(x, dim_target)

Shape:

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

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

Parameters

• x (Tensor) – an input tensor.

• dim_target (int) – typically -2 or -3.

Returns

x_aggregated

Return type

Tensor

class weavenet.layers.DualSoftmax(dim_src=-3, dim_tar=-2)

Applies the dual-softmax calculation to a batched matrices. DualSoftMax is originally proposed in LoFTR (CVPR2021).

\[\text{DualSoftmax}(x^{ab}_{ij}, x^{ba}_{ij}) = \frac{\exp(x^{ab}_{ij})}{\sum_j \exp(x^{ab}_{ij})} * \frac{\exp(x^{ba}_{ij})}{\sum_i \exp(x^{ba}_{ij})}\]

In original definition, always \({x^{ba}}^T=x^{ab}\). This is an extensional implementation that accepts \(x^{ba}\neq x^{ab}\) to input the two stream outputs of WeaveNet.

Parameters

• dim_src (int) –

• dim_tar (int) –

forward(xab, xba_t=None)

Calculate the dual softmax for batched matrices.

Shape:

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

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

• output: \((B, \ldots, N, M, 1)\)

Parameters

• xab (Tensor) – 1st batched matrices.

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

Returns

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

Return type

Tuple[Tensor, Tensor, Tensor]

class weavenet.layers.DualSoftmaxSqrt(dim_src=-3, dim_tar=-2)

A variation of DualSoftmax. This variation is effective when xab and xba derive from different computational flow since the forward sqrt operation amplify the backward gradient to each flow.

\[\text{DualSoftmaxSqrt}(x^{ab}_{ij}, x^{ba}_{ij}) = \sqrt{\text{DualSoftmax}(x^{ab}_{ij}, x^{ba}_{ij})}\]

Parameters

• dim_src (int) –

• dim_tar (int) –

forward(xab, xba_t=None)

Shape and Args: same as DualSoftmax

Parameters

• xab (Tensor) – 1st batched matrices.

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

Returns

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

Return type

Tuple[Tensor, Tensor, Tensor]

class weavenet.layers.BatchNormXXC(C)

Applies BatchNorm1d to \((\ldots, C)\)-shaped tensors. This module is prepered since nn.BatchNorm2d assumes the input format of \((B, C, H, W)\) but if kernel size is 1, nn.Conv2d to \((B, C, H, W)\) is slower than nn.Linear to \((B, H, W, C)\), which is our case for bipartite-graph edge embedding of \((B, N, M, C)\). In addition, nn.Linear works well with sparse cases.

Example of Usage:

# assume batch_size=8, the problem instance size is 5x5,  and each edge feature is 32 channels.
B, N, M, C = 8, 5, 5, 32
linear = nn.Linear(32, 64)
bn = BatchNormXXC(64)
x = torch.rand((B, N, M, C), dtype=torch.float) # prepare a random input.
x = linear(x)
x = bn(x)

forward(x)

Shape:

• x: \((\ldots, C)\)

• output: \((\ldots, C)\)

Parameters

x (Tensor) – target variable.

Returns

batch-normed variable.

Return type

x

Experimental

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

Applies a variation of SetEncoderPointNet. This class max-pools in dim_src direction in addition to dim_target direction of standard SetEncoder.

Parameters

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

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

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

forward(x, dim_target)

Shape and Args: same as SetEncoderPointNet

Parameters

• x (Tensor) – an input tensor of edge features.

• dim_target (int) – the dimension of edge’s target vertex and must be -3 or -2.

Returns

z_edge_features

Return type

Tensor

class weavenet.layers.SetEncoderPointNetTotalDirectional(first_process, aggregator, second_process_edge, second_process_vertex)

Applies a variation of SetEncoderPointNet. This class max-pools all the edge features in addition to SetEncoderPointNetCrossDirectional, and concatenate the summerized features to original edge feature.

Parameters

• in_channels – the number of input channels.

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

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

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

• aggregator (Callable[[Tensor, int], Tensor]) –

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

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

forward(x, dim_target)

Shape:

• x: \((\ldots)\) (not defined with this abstractive class)

• output: \((\ldots)\)　 (not defined with this abstractive class)

Parameters

• x (Tensor) – an input tensor.

• dim_target (int) –

Returns

z_edge_features, z_vertex_features

Return type

Tensor

class weavenet.layers.CrossDifferenceConcat(dim_feature=-1)

Applies cross-concatenation of mean and difference (experimental).

\[\text{CrossDiffConcat}([x^{ab}, {x^{ba}}^\top]) = [\text{cat}([x^{ab}, {x^{ba}}^\top - x^{ab}], \text{dim}=-1), \text{cat}([{x^{ba}}^\top,x^{ab}-{x^{ba}}^\top], \text{dim}=-1)]\]

Note

This class was not very effective with stable matching test.

Parameters

dim_feature (int) –

output_channels(input_channels)

Parameters

input_channels (int) – assumed input channels.

Returns

\(2*\) input_channels

Return type

output_channels

forward(xab, xba_t)

Shape:

• xab: \((\ldots, C)\)

• xba_t: \((\ldots, C)\)

• output(zab, zba_t): \([(\ldots, 2*C),(\ldots, 2*C)]\)

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

Tuple[zab, zba_t] calculated as \(\text{CrossDiffConcat}([x^{ab}, {x^{ba}}^\top])\) .

Return type

Tuple[Tensor, Tensor]

class weavenet.layers.DualSoftmaxFuzzyLogicAnd(dim_src=-3, dim_tar=-2)

Applies the calculation proposed in Shira Li, “Deep Learning for Two-Sided Matching Markets”.

\[\text{DualSoftmaxFuzzyLogicAnd}(x^{ab}_{ij}, x^{ba}_{ij}) = \min(\frac{\exp(x^{ab}_{ij})}{\sum_j \exp(x^{ab}_{ij})}, \frac{\exp(x^{ba}_{ij})}{\sum_i \exp(x^{ba}_{ij})})\]

Parameters

• dim_src (int) –

• dim_tar (int) –

forward(xab, xba_t=None)

Shape and Args: same as DualSoftmax

Parameters

• xab (Tensor) – 1st batched matrices.

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

Returns

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

Return type

Tuple[Tensor, Tensor, Tensor]

Private Layers

class weavenet.layers.SetEncoderBase(first_process, aggregator, second_process_edge, second_process_vertex)

Applies abstracted set-encoding process.

\[\text{SetEncoderBase}(x) = \text{second_process}(\text{merger}(x, \text{aggregator}(\text{first_process}(x))))\]

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, int], 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, dim_target)

Shape:

• x: \((\ldots, N, M, \text{in_channels})\)

• output: \((\ldots, N, M, \text{output_channels})\)

Parameters

• x (Tensor) – an input tensor.

• dim_target (int) – the dimension of edge’s target vertex.

Returns

z_edge_features

Return type

Tensor

class weavenet.layers.Interactor

Abstract CrossConcat and any other interactor between feature blocks of two stream architecture. It must have a function output_channels() to report its resultant feature’s output channels (estimated based on the input_channels).

output_channels(input_channels)

Parameters

input_channels (int) – assumed input channels.

Returns

input_channels (dummy)

Return type

output_channels

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

Tuple[Tensor, Tensor]