equiadapt.pointcloud.canonicalization_networks package

Submodules

equiadapt.pointcloud.canonicalization_networks.equivariant_networks module

class equiadapt.pointcloud.canonicalization_networks.equivariant_networks.VNSmall(hyperparams: DictConfig)[source]

Bases: Module

VNSmall is a small variant of the vector neuron equivariant network used for canonicalization of point clouds.

Parameters:: hyperparams (DictConfig) – Hyperparameters for the network.

n_knn

Number of nearest neighbors to consider.

Type:: int

pooling

Pooling type to use, either “max” or “mean”.

Type:: str

conv_pos

Convolutional layer for positional encoding.

Type:: VNLinearLeakyReLU

conv1

First convolutional layer.

Type:: VNLinearLeakyReLU

bn1

Batch normalization layer.

Type:: VNBatchNorm

conv2

Second convolutional layer.

Type:: VNLinearLeakyReLU

dropout

Dropout layer.

Type:: nn.Dropout

pool

Pooling layer.

Type:: Union[VNMaxPool, mean_pool]

__init__()[source]: Initializes the VNSmall network.

forward()[source]: Forward pass of the VNSmall network.

forward(point_cloud: Tensor) → Tensor[source]

Forward pass of the VNSmall network.

For every pointcloud in the batch, the network outputs three vectors that transform equivariantly with respect to SO3 group.

Parameters:: point_cloud (torch.Tensor) – Input point cloud tensor of shape (batch_size, num_points, 3).
Returns:: Output tensor of shape (batch_size, 3, 3).
Return type:: torch.Tensor

equiadapt.pointcloud.canonicalization_networks.equivariant_networks.get_graph_feature_cross(x: Tensor, k: int = 20, idx: Tensor | None = None) → Tensor[source]

Computes the graph feature cross for a given input tensor.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_size, num_dims, num_points).
k (int, optional) – The number of nearest neighbors to consider. Defaults to 20.
idx (torch.Tensor, optional) – The indices of the nearest neighbors. Defaults to None.

Returns:

The computed graph feature cross tensor of shape (batch_size, num_dims*3, num_points, k).

Return type:

torch.Tensor

equiadapt.pointcloud.canonicalization_networks.equivariant_networks.knn(x: Tensor, k: int) → Tensor[source]

Performs k-nearest neighbors search on a given set of points.

Parameters:

x (torch.Tensor) – The input tensor representing a set of points. Shape: (batch_size, num_points, num_dimensions).
k (int) – The number of nearest neighbors to find.

Returns:

The indices of the k nearest neighbors for each point in x.: Shape: (batch_size, num_points, k).

Return type:

torch.Tensor

equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers module

Layers for vector neuron networks

Taken from Vector Neurons: A General Framework for SO(3)-Equivariant Networks (https://arxiv.org/abs/2104.12229) paper and their codebase https://github.com/FlyingGiraffe/vnn

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNBatchNorm(num_features: int, dim: int)[source]

Bases: Module

Vector Neuron Batch Normalization layer.

VNBatchNorm applies batch normalization to the input features.

__init__()[source]: Initializes the VNBatchNorm layer.

forward()[source]: Performs forward pass of the VNBatchNorm layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the Vector Neuron Batch Normalization layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor after applying batch normalization.
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNBilinear(in_channels1: int, in_channels2: int, out_channels: int)[source]

Bases: Module

Vector Neuron Bilinear layer.

VNBilinear applies a bilinear layer to the input features.

__init__()[source]: Initializes the VNBilinear layer.

forward()[source]: Performs forward pass of the VNBilinear layer.

forward(x: Tensor, labels: Tensor) → Tensor[source]

Forward pass of the VNBilinear layer.

Parameters:

x (torch.Tensor) – Input features of shape [B, N_feat, 3, N_samples, …].
labels (torch.Tensor) – Labels of shape [B, N_feat, N_samples].

Returns:

Output features after applying the bilinear transformation.

Return type:

torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNLeakyReLU(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Vector Neuron Leaky ReLU layer.

VNLLeakyReLU applies a LeakyReLU activation to the input features.

__init__()[source]: Initializes the VNLeakyReLU layer.

forward()[source]: Performs forward pass of the VNLeakyReLU layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the VNLeakyReLU module.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor after applying VNLeakyReLU activation.
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNLinear(in_channels: int, out_channels: int)[source]

Bases: Module

Vector Neuron Linear layer.

This layer applies a linear transformation to the input tensor.

__init__()[source]: Initializes the VNLinear layer.

forward()[source]: Performs forward pass of the VNLinear layer.

forward(x: Tensor) → Tensor[source]

Performs forward pass of the VNLinear layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor of shape [B, N_feat, 3, N_samples, …].
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNLinearLeakyReLU(in_channels: int, out_channels: int, dim: int = 5, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Vector Neuron Linear Leaky ReLU layer.

VNLinearLeakyReLU applies a linear transformation followed by a LeakyReLU activation to the input features.

__init__()[source]: Initializes the VNLinearLeakyReLU layer.

forward()[source]: Performs forward pass of the VNLinearLeakyReLU layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the VNLinearLeakyReLU layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …]
Returns:: Output tensor of shape [B, N_feat, 3, N_samples, …]
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNMaxPool(in_channels: int)[source]

Bases: Module

Vector Neuron Max Pooling layer.

VNMaxPool applies max pooling to the input features.

__init__()[source]: Initializes the VNMaxPool layer.

forward()[source]: Performs forward pass of the VNMaxPool layer.

forward(x: Tensor) → Tensor[source]

Performs vector neuron max pooling on the input tensor.

Parameters:: x (torch.Tensor) – Point features of shape [B, N_feat, 3, N_samples, …].
Returns:: Max pooled tensor of shape [B, N_feat, 3, N_samples, …].
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNSoftplus(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.0)[source]

Bases: Module

Vector Neuron Softplus layer.

VNSoftplus applies a softplus activation to the input features.

__init__()[source]: Initializes the VNSoftplus layer.

forward()[source]: Performs forward pass of the VNSoftplus layer.

forward(x: Tensor) → Tensor[source]

Performs forward pass of the VNSoftplus layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor of shape [B, N_feat, 3, N_samples, …].
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.VNStdFeature(in_channels: int, dim: int = 4, normalize_frame: bool = False, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Vector Neuron Standard Feature module.

This module performs standard feature extraction using Vector Neuron layers. It takes point features as input and applies a series of VNLinearLeakyReLU layers followed by a linear layer to produce the standard features.

dim

Dimension of the input features.

Type:: int

normalize_frame

Whether to normalize the frame.

Type:: bool

__init__()[source]: Initializes the VNStdFeature module.

forward()[source]: Performs forward pass of the VNStdFeature module.

Shape:

Input: (B, N_feat, 3, N_samples, …)
Output:
- x_std: (B, N_feat, dim, N_samples, …)
- z0: (B, dim, 3)

Example

>>> model = VNStdFeature(in_channels=64, dim=4, normalize_frame=True)
>>> input = torch.randn(2, 64, 3, 100)
>>> output, frame_vectors = model(input)

forward(x: Tensor) → Tuple[Tensor, Tensor][source]

Forward pass of the VNStdFeature module.

Parameters:: x (torch.Tensor) – Input point features of shape (B, N_feat, 3, N_samples, …).
Returns:: Tuple containing the standard features and the frame vectors.
Return type:: Tuple[torch.Tensor, torch.Tensor]

Note

The frame vectors are computed only if normalize_frame is set to True.
The shape of the standard features depends on the value of dim.

equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers.mean_pool(x: Tensor, dim: int = -1, keepdim: bool = False) → Tensor[source]

Compute the mean pooling of a tensor along a specified dimension.

Parameters:

x (torch.Tensor) – The input tensor.
dim (int, optional) – The dimension along which to compute the mean pooling. Default is -1.
keepdim (bool, optional) – Whether to keep the dimension of the input tensor. Default is False.

Returns:

The mean pooled tensor.

Return type:

torch.Tensor

Module contents

This package contains equivariant modules and networks for the equiadapt pointcloud canonicalization.

class equiadapt.pointcloud.canonicalization_networks.VNBatchNorm(num_features: int, dim: int)[source]

Bases: Module

Vector Neuron Batch Normalization layer.

VNBatchNorm applies batch normalization to the input features.

__init__()[source]: Initializes the VNBatchNorm layer.

forward()[source]: Performs forward pass of the VNBatchNorm layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the Vector Neuron Batch Normalization layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor after applying batch normalization.
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNBilinear(in_channels1: int, in_channels2: int, out_channels: int)[source]

Bases: Module

Vector Neuron Bilinear layer.

VNBilinear applies a bilinear layer to the input features.

__init__()[source]: Initializes the VNBilinear layer.

forward()[source]: Performs forward pass of the VNBilinear layer.

forward(x: Tensor, labels: Tensor) → Tensor[source]

Forward pass of the VNBilinear layer.

Parameters:

x (torch.Tensor) – Input features of shape [B, N_feat, 3, N_samples, …].
labels (torch.Tensor) – Labels of shape [B, N_feat, N_samples].

Returns:

Output features after applying the bilinear transformation.

Return type:

torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNLeakyReLU(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Vector Neuron Leaky ReLU layer.

VNLLeakyReLU applies a LeakyReLU activation to the input features.

__init__()[source]: Initializes the VNLeakyReLU layer.

forward()[source]: Performs forward pass of the VNLeakyReLU layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the VNLeakyReLU module.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor after applying VNLeakyReLU activation.
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNLinear(in_channels: int, out_channels: int)[source]

Bases: Module

Vector Neuron Linear layer.

This layer applies a linear transformation to the input tensor.

__init__()[source]: Initializes the VNLinear layer.

forward()[source]: Performs forward pass of the VNLinear layer.

forward(x: Tensor) → Tensor[source]

Performs forward pass of the VNLinear layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor of shape [B, N_feat, 3, N_samples, …].
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNLinearLeakyReLU(in_channels: int, out_channels: int, dim: int = 5, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Vector Neuron Linear Leaky ReLU layer.

VNLinearLeakyReLU applies a linear transformation followed by a LeakyReLU activation to the input features.

__init__()[source]: Initializes the VNLinearLeakyReLU layer.

forward()[source]: Performs forward pass of the VNLinearLeakyReLU layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the VNLinearLeakyReLU layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …]
Returns:: Output tensor of shape [B, N_feat, 3, N_samples, …]
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNMaxPool(in_channels: int)[source]

Bases: Module

Vector Neuron Max Pooling layer.

VNMaxPool applies max pooling to the input features.

__init__()[source]: Initializes the VNMaxPool layer.

forward()[source]: Performs forward pass of the VNMaxPool layer.

forward(x: Tensor) → Tensor[source]

Performs vector neuron max pooling on the input tensor.

Parameters:: x (torch.Tensor) – Point features of shape [B, N_feat, 3, N_samples, …].
Returns:: Max pooled tensor of shape [B, N_feat, 3, N_samples, …].
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNSmall(hyperparams: DictConfig)[source]

Bases: Module

VNSmall is a small variant of the vector neuron equivariant network used for canonicalization of point clouds.

Parameters:: hyperparams (DictConfig) – Hyperparameters for the network.

n_knn

Number of nearest neighbors to consider.

Type:: int

pooling

Pooling type to use, either “max” or “mean”.

Type:: str

conv_pos

Convolutional layer for positional encoding.

Type:: VNLinearLeakyReLU

conv1

First convolutional layer.

Type:: VNLinearLeakyReLU

bn1

Batch normalization layer.

Type:: VNBatchNorm

conv2

Second convolutional layer.

Type:: VNLinearLeakyReLU

dropout

Dropout layer.

Type:: nn.Dropout

pool

Pooling layer.

Type:: Union[VNMaxPool, mean_pool]

__init__()[source]: Initializes the VNSmall network.

forward()[source]: Forward pass of the VNSmall network.

forward(point_cloud: Tensor) → Tensor[source]

Forward pass of the VNSmall network.

For every pointcloud in the batch, the network outputs three vectors that transform equivariantly with respect to SO3 group.

Parameters:: point_cloud (torch.Tensor) – Input point cloud tensor of shape (batch_size, num_points, 3).
Returns:: Output tensor of shape (batch_size, 3, 3).
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNSoftplus(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.0)[source]

Bases: Module

Vector Neuron Softplus layer.

VNSoftplus applies a softplus activation to the input features.

__init__()[source]: Initializes the VNSoftplus layer.

forward()[source]: Performs forward pass of the VNSoftplus layer.

forward(x: Tensor) → Tensor[source]

Performs forward pass of the VNSoftplus layer.

Parameters:: x (torch.Tensor) – Input tensor of shape [B, N_feat, 3, N_samples, …].
Returns:: Output tensor of shape [B, N_feat, 3, N_samples, …].
Return type:: torch.Tensor

class equiadapt.pointcloud.canonicalization_networks.VNStdFeature(in_channels: int, dim: int = 4, normalize_frame: bool = False, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Vector Neuron Standard Feature module.

This module performs standard feature extraction using Vector Neuron layers. It takes point features as input and applies a series of VNLinearLeakyReLU layers followed by a linear layer to produce the standard features.

dim

Dimension of the input features.

Type:: int

normalize_frame

Whether to normalize the frame.

Type:: bool

__init__()[source]: Initializes the VNStdFeature module.

forward()[source]: Performs forward pass of the VNStdFeature module.

Shape:

Input: (B, N_feat, 3, N_samples, …)
Output:
- x_std: (B, N_feat, dim, N_samples, …)
- z0: (B, dim, 3)

Example

>>> model = VNStdFeature(in_channels=64, dim=4, normalize_frame=True)
>>> input = torch.randn(2, 64, 3, 100)
>>> output, frame_vectors = model(input)

forward(x: Tensor) → Tuple[Tensor, Tensor][source]

Forward pass of the VNStdFeature module.

Parameters:: x (torch.Tensor) – Input point features of shape (B, N_feat, 3, N_samples, …).
Returns:: Tuple containing the standard features and the frame vectors.
Return type:: Tuple[torch.Tensor, torch.Tensor]

Note

The frame vectors are computed only if normalize_frame is set to True.
The shape of the standard features depends on the value of dim.

equiadapt.pointcloud.canonicalization_networks.get_graph_feature_cross(x: Tensor, k: int = 20, idx: Tensor | None = None) → Tensor[source]

Computes the graph feature cross for a given input tensor.

Parameters:

x (torch.Tensor) – The input tensor of shape (batch_size, num_dims, num_points).
k (int, optional) – The number of nearest neighbors to consider. Defaults to 20.
idx (torch.Tensor, optional) – The indices of the nearest neighbors. Defaults to None.

Returns:

The computed graph feature cross tensor of shape (batch_size, num_dims*3, num_points, k).

Return type:

torch.Tensor

equiadapt.pointcloud.canonicalization_networks.knn(x: Tensor, k: int) → Tensor[source]

Performs k-nearest neighbors search on a given set of points.

Parameters:

x (torch.Tensor) – The input tensor representing a set of points. Shape: (batch_size, num_points, num_dimensions).
k (int) – The number of nearest neighbors to find.

Returns:

The indices of the k nearest neighbors for each point in x.: Shape: (batch_size, num_points, k).

Return type:

torch.Tensor

equiadapt.pointcloud.canonicalization_networks.mean_pool(x: Tensor, dim: int = -1, keepdim: bool = False) → Tensor[source]

Compute the mean pooling of a tensor along a specified dimension.

Parameters:

x (torch.Tensor) – The input tensor.
dim (int, optional) – The dimension along which to compute the mean pooling. Default is -1.
keepdim (bool, optional) – Whether to keep the dimension of the input tensor. Default is False.

Returns:

The mean pooled tensor.

Return type:

torch.Tensor