equiadapt.pointcloud package

Subpackages

• equiadapt.pointcloud.canonicalization package
  • Submodules
  • equiadapt.pointcloud.canonicalization.continuous_group module
    • ContinuousGroupPointcloudCanonicalization
      • ContinuousGroupPointcloudCanonicalization.device
      • ContinuousGroupPointcloudCanonicalization.get_groupelement()
      • ContinuousGroupPointcloudCanonicalization.canonicalize()
      • ContinuousGroupPointcloudCanonicalization.canonicalize()
      • ContinuousGroupPointcloudCanonicalization.get_groupelement()
    • EquivariantPointcloudCanonicalization
      • EquivariantPointcloudCanonicalization.canonicalization_network
      • EquivariantPointcloudCanonicalization.canonicalization_hyperparams
      • EquivariantPointcloudCanonicalization.canonicalization_info_dict
      • EquivariantPointcloudCanonicalization.get_groupelement()
  • Module contents
    • ContinuousGroupPointcloudCanonicalization
      • ContinuousGroupPointcloudCanonicalization.device
      • ContinuousGroupPointcloudCanonicalization.get_groupelement()
      • ContinuousGroupPointcloudCanonicalization.canonicalize()
      • ContinuousGroupPointcloudCanonicalization.canonicalize()
      • ContinuousGroupPointcloudCanonicalization.get_groupelement()
    • EquivariantPointcloudCanonicalization
      • EquivariantPointcloudCanonicalization.canonicalization_network
      • EquivariantPointcloudCanonicalization.canonicalization_hyperparams
      • EquivariantPointcloudCanonicalization.canonicalization_info_dict
      • EquivariantPointcloudCanonicalization.get_groupelement()
• equiadapt.pointcloud.canonicalization_networks package
  • Submodules
  • equiadapt.pointcloud.canonicalization_networks.equivariant_networks module
    • VNSmall
      • VNSmall.n_knn
      • VNSmall.pooling
      • VNSmall.conv_pos
      • VNSmall.conv1
      • VNSmall.bn1
      • VNSmall.conv2
      • VNSmall.dropout
      • VNSmall.pool
      • VNSmall.__init__()
      • VNSmall.forward()
      • VNSmall.forward()
    • get_graph_feature_cross()
    • knn()
  • equiadapt.pointcloud.canonicalization_networks.vector_neuron_layers module
    • VNBatchNorm
      • VNBatchNorm.__init__()
      • VNBatchNorm.forward()
      • VNBatchNorm.forward()
    • VNBilinear
      • VNBilinear.__init__()
      • VNBilinear.forward()
      • VNBilinear.forward()
    • VNLeakyReLU
      • VNLeakyReLU.__init__()
      • VNLeakyReLU.forward()
      • VNLeakyReLU.forward()
    • VNLinear
      • VNLinear.__init__()
      • VNLinear.forward()
      • VNLinear.forward()
    • VNLinearLeakyReLU
      • VNLinearLeakyReLU.__init__()
      • VNLinearLeakyReLU.forward()
      • VNLinearLeakyReLU.forward()
    • VNMaxPool
      • VNMaxPool.__init__()
      • VNMaxPool.forward()
      • VNMaxPool.forward()
    • VNSoftplus
      • VNSoftplus.__init__()
      • VNSoftplus.forward()
      • VNSoftplus.forward()
    • VNStdFeature
      • VNStdFeature.dim
      • VNStdFeature.normalize_frame
      • VNStdFeature.__init__()
      • VNStdFeature.forward()
      • VNStdFeature.forward()
    • mean_pool()
  • Module contents
    • VNBatchNorm
      • VNBatchNorm.__init__()
      • VNBatchNorm.forward()
      • VNBatchNorm.forward()
    • VNBilinear
      • VNBilinear.__init__()
      • VNBilinear.forward()
      • VNBilinear.forward()
    • VNLeakyReLU
      • VNLeakyReLU.__init__()
      • VNLeakyReLU.forward()
      • VNLeakyReLU.forward()
    • VNLinear
      • VNLinear.__init__()
      • VNLinear.forward()
      • VNLinear.forward()
    • VNLinearLeakyReLU
      • VNLinearLeakyReLU.__init__()
      • VNLinearLeakyReLU.forward()
      • VNLinearLeakyReLU.forward()
    • VNMaxPool
      • VNMaxPool.__init__()
      • VNMaxPool.forward()
      • VNMaxPool.forward()
    • VNSmall
      • VNSmall.n_knn
      • VNSmall.pooling
      • VNSmall.conv_pos
      • VNSmall.conv1
      • VNSmall.bn1
      • VNSmall.conv2
      • VNSmall.dropout
      • VNSmall.pool
      • VNSmall.__init__()
      • VNSmall.forward()
      • VNSmall.forward()
    • VNSoftplus
      • VNSoftplus.__init__()
      • VNSoftplus.forward()
      • VNSoftplus.forward()
    • VNStdFeature
      • VNStdFeature.dim
      • VNStdFeature.normalize_frame
      • VNStdFeature.__init__()
      • VNStdFeature.forward()
      • VNStdFeature.forward()
    • get_graph_feature_cross()
    • knn()
    • mean_pool()

Module contents

This package contains modules for the equiadapt pointcloud canonicalization.

class equiadapt.pointcloud.ContinuousGroupPointcloudCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig)

Bases: ContinuousGroupCanonicalization

This class represents a continuous group point cloud canonicalization.

Parameters:

• canonicalization_network (torch.nn.Module) – The canonicalization network.

• canonicalization_hyperparams (DictConfig) – The hyperparameters for canonicalization.

device

The device on which the operations are performed.

get_groupelement()

Maps the input point cloud to the group element.

canonicalize()

Returns the canonicalized point cloud.

canonicalize(x: Tensor, targets: List | None = None, **kwargs: Any) → Tensor | Tuple[Tensor, List]

This method takes an image as input and returns the canonicalized image.

Parameters:

• x (torch.Tensor) – The input point cloud.

• targets (Optional[List]) – The list of targets (optional).

• **kwargs (Any) – Additional keyword arguments.

Returns:

The canonicalized point cloud.

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, List]]

get_groupelement(x: Tensor) → dict

This method takes the input image and maps it to the group element.

Parameters:

x (torch.Tensor) – The input image.

Returns:

The group element.

Return type:

dict

Raises:

NotImplementedError – If the method is not implemented.

class equiadapt.pointcloud.EquivariantPointcloudCanonicalization(canonicalization_network: Module, canonicalization_hyperparams: DictConfig)

Bases: ContinuousGroupPointcloudCanonicalization

This class represents the equivariant point cloud canonicalization module.

It inherits from the ContinuousGroupPointcloudCanonicalization class.

Parameters:

• canonicalization_network (torch.nn.Module) – The canonicalization network module.

• canonicalization_hyperparams (DictConfig) – The hyperparameters for the canonicalization.

canonicalization_network

The canonicalization network module.

Type:

torch.nn.Module

canonicalization_hyperparams

The hyperparameters for the canonicalization.

Type:

DictConfig

canonicalization_info_dict

A dictionary to store the canonicalization information.

Type:

dict

get_groupelement(x: Tensor) → Dict[str, Tensor]

This method takes the input image and maps it to the group element.

Parameters:

x (torch.Tensor) – The input point cloud.

Returns:

A dictionary containing the group element.

Return type:

Dict[str, torch.Tensor]

class equiadapt.pointcloud.VNBatchNorm(num_features: int, dim: int)

Bases: Module

Vector Neuron Batch Normalization layer.

VNBatchNorm applies batch normalization to the input features.

__init__()

Initializes the VNBatchNorm layer.

forward()

Performs forward pass of the VNBatchNorm layer.

forward(x: Tensor) → Tensor

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.VNBilinear(in_channels1: int, in_channels2: int, out_channels: int)

Bases: Module

Vector Neuron Bilinear layer.

VNBilinear applies a bilinear layer to the input features.

__init__()

Initializes the VNBilinear layer.

forward()

Performs forward pass of the VNBilinear layer.

forward(x: Tensor, labels: Tensor) → Tensor

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.VNLeakyReLU(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.2)

Bases: Module

Vector Neuron Leaky ReLU layer.

VNLLeakyReLU applies a LeakyReLU activation to the input features.

__init__()

Initializes the VNLeakyReLU layer.

forward()

Performs forward pass of the VNLeakyReLU layer.

forward(x: Tensor) → Tensor

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.VNLinear(in_channels: int, out_channels: int)

Bases: Module

Vector Neuron Linear layer.

This layer applies a linear transformation to the input tensor.

__init__()

Initializes the VNLinear layer.

forward()

Performs forward pass of the VNLinear layer.

forward(x: Tensor) → Tensor

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.VNLinearLeakyReLU(in_channels: int, out_channels: int, dim: int = 5, share_nonlinearity: bool = False, negative_slope: float = 0.2)

Bases: Module

Vector Neuron Linear Leaky ReLU layer.

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

__init__()

Initializes the VNLinearLeakyReLU layer.

forward()

Performs forward pass of the VNLinearLeakyReLU layer.

forward(x: Tensor) → Tensor

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.VNMaxPool(in_channels: int)

Bases: Module

Vector Neuron Max Pooling layer.

VNMaxPool applies max pooling to the input features.

__init__()

Initializes the VNMaxPool layer.

forward()

Performs forward pass of the VNMaxPool layer.

forward(x: Tensor) → Tensor

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.VNSmall(hyperparams: DictConfig)

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__()

Initializes the VNSmall network.

forward()

Forward pass of the VNSmall network.

forward(point_cloud: Tensor) → Tensor

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.VNSoftplus(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.0)

Bases: Module

Vector Neuron Softplus layer.

VNSoftplus applies a softplus activation to the input features.

__init__()

Initializes the VNSoftplus layer.

forward()

Performs forward pass of the VNSoftplus layer.

forward(x: Tensor) → Tensor

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.VNStdFeature(in_channels: int, dim: int = 4, normalize_frame: bool = False, share_nonlinearity: bool = False, negative_slope: float = 0.2)

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__()

Initializes the VNStdFeature module.

forward()

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]

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.get_graph_feature_cross(x: Tensor, k: int = 20, idx: Tensor | None = None) → Tensor

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.knn(x: Tensor, k: int) → Tensor

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.mean_pool(x: Tensor, dim: int = -1, keepdim: bool = False) → Tensor

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