Class VNInFeature

src/shape_assembly/models/encoder/vn_layers.py:240–297

VN-Invariant layer.

class VNInFeature(nn.Module):
    """VN-Invariant layer."""

    def __init__(
            self,
            in_channels,
            dim=4,
            share_nonlinearity=False,
            negative_slope=0.2,
            use_rmat=False,
    ):
        super().__init__()

        self.dim = dim
        self.use_rmat = use_rmat
        self.vn1 = VNLinearBNLeakyReLU(
            in_channels,
            in_channels // 2,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            negative_slope=negative_slope,
        )
        self.vn2 = VNLinearBNLeakyReLU(
            in_channels // 2,
            in_channels // 4,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            negative_slope=negative_slope,
        )
        self.vn_lin = conv1x1(
            in_channels // 4, 2 if self.use_rmat else 3, dim=dim)

    def forward(self, x):
        """
        Args:
            x: point features of shape [B, C, 3, N, ...]
        Returns:
            rotation invariant features of the same shape
        """
        z = self.vn1(x)
        z = self.vn2(z)
        z = self.vn_lin(z)  # [B, 3, 3, N] or [B, 2, 3, N]
        if self.use_rmat:
            z = z.flatten(1, 2).transpose(1, 2).contiguous()  # [B, N, 6]
            z = rot6d_to_matrix(z)  # [B, N, 3, 3]
            z = z.permute(0, 2, 3, 1)  # [B, 3, 3, N]
        z = z.transpose(1, 2).contiguous()

        if self.dim == 4:
            x_in = torch.einsum('bijm,bjkm->bikm', x, z)
        elif self.dim == 3:
            x_in = torch.einsum('bij,bjk->bik', x, z)
        elif self.dim == 5:
            x_in = torch.einsum('bijmn,bjkmn->bikmn', x, z)
        else:
            raise NotImplementedError(f'dim={self.dim} is not supported')

        return x_in