aidsorb.modules

torch.nn.Module’s for point cloud processing.

Note

PointNetBackbone, PointNetClsHead and PointNetSegHead have their initial layers lazy initialized, so you don’t need to specify the input dimensionality.

Warning

It is recommended to use batched inputs in all cases. For example, even if a single pcd of shape (3+C, N) is to be passed to PointNetBackbone, reshape it to (1, 3+C, N). One way to you can do it is the following: pcd = pcd.unsqueeze(0).

Todo

Add more modules for point cloud processing.

References

[PointNet] (1,2,3,4,5)

R. Q. Charles, H. Su, M. Kaichun and L. J. Guibas, “PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation,” 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Honolulu, HI, USA, 2017, pp. 77-85, doi: 10.1109/CVPR.2017.16.

class aidsorb.modules.PointNet(head, n_global_feats=1024, local_feats=False)

Bases: Module

Vanilla version from the [PointNet] paper where TNet’s have been removed.

PointNet takes as input a point cloud and produces one or more outputs. The type of the task is determined by head.

Currently implemented heads include:

1. PointNetClsHead: classification and regression

2. PointNetSegHead: segmentation

The input must be batched, i.e. have shape of (B, C, N) where B is the batch size, C is the number of input channels and N is the number of points in each point cloud.

Tip

You can define a custom_head head as a torch.nn.Module and pass it to head.

If local_features=False, the input to custom_head must have the same shape as in PointNetClsHead.forward(). Otherwise, the input to custom_head must have the same shape as in PointNetSegHead.forward().

Parameters:

• head (torch.nn.Module)

• n_global_feats (int, default=1024)

• local_feats (bool, default=False)

See also

modules.PointNetBackbone

For a description of local_feats and n_global_feats.

Examples

>>> cls_head = PointNetClsHead(n_outputs=2)
>>> seg_head = PointNetSegHead(n_outputs=10)
>>> x = torch.randn(32, 4, 300)

>>> cls_net = PointNet(cls_head, 256)
>>> cls_net(x).shape
torch.Size([32, 2])
>>> cls_net.backbone(x).shape  # Only features.
torch.Size([32, 256])

>>> seg_net = PointNet(head=seg_head, n_global_feats=512, local_feats=True)
>>> seg_net(x).shape
torch.Size([32, 300, 10])
>>> seg_net.backbone(x, True)[1].shape  # Features and critical indices.
torch.Size([32, 512])

forward(x)

Run the forward pass.

Parameters:

x (tensor of shape (B, C, N))

Returns:

out – Output of head.

Return type:

tensor

class aidsorb.modules.PointNetBackbone(n_global_feats=1024, local_feats=False)

Bases: Module

Backbone of the vanilla version from the [PointNet] paper, where TNet’s have been removed.

This module extracts features which can then be passed to a task head for predictions. This module also returns the critical indices.

The input must be batched, i.e. have shape of (B, C, N) where B is the batch size, C is the number of input channels and N is the number of points in each point cloud.

Parameters:

• n_global_feats (int, default=1024) – Number of global features.

• local_feats (bool, default=False) – If True, the returned features are a concatenation of local features and global features. Otherwise, the global features are returned.

Examples

>>> feat = PointNetBackbone(2048)
>>> x = torch.randn(32, 4, 200)
>>> features, indices = feat(x, return_indices=True)
>>> features.shape
torch.Size([32, 2048])
>>> indices.shape
torch.Size([32, 2048])

>>> feat = PointNetBackbone(1024, True)
>>> x = torch.randn(16, 4, 100)
>>> features, indices = feat(x, return_indices=True)
>>> features.shape
torch.Size([16, 1088, 100])
>>> indices.shape
torch.Size([16, 1024])

>>> feat =  PointNetBackbone(512)
>>> x = torch.randn(8, 3, 50)
>>> feat(x).shape  # Only features, no critical indices.
torch.Size([8, 512])

forward(x, return_indices=False)

Return the features and optionally critical indices.

The type of the features is determined by local_feats.

Parameters:

• x (tensor of shape (B, C, N))

• return_indices (bool, default=False) – Whether to return critical indices.

Returns:

out – If return_indices=False the output are the features, otherwise tuple of the form (features, critical_indices).

Return type:

tensor or tuple of tensors

class aidsorb.modules.PointNetClsHead(n_outputs=1, dropout_rate=0)

Bases: Module

Classification head from the [PointNet] paper.

Tip

This head can be used for classification or regression tasks.

Parameters:

• n_outputs (int, default=1)

• dropout_rate (float, default=0)

Examples

>>> head = PointNetClsHead(n_outputs=4)
>>> x = torch.randn(64, 13)
>>> head(x).shape
torch.Size([64, 4])

forward(x)

Run the forward pass.

Parameters:

x (tensor of shape (B, C))

Returns:

out

Return type:

tensor of shape (B, n_outputs)

class aidsorb.modules.PointNetSegHead(n_outputs=1)

Bases: Module

Segmentation head from the [PointNet] paper.

Tip

This head can be used for segmentation tasks.

Parameters:

n_outputs (int, default=1)

Examples

>>> head = PointNetSegHead(n_outputs=2)
>>> x = torch.randn(32, 1088, 400)
>>> head(x).shape
torch.Size([32, 400, 2])

forward(x)

Run the forward pass.

Parameters:

x (tensor of shape (B, C, N).)

Returns:

out

Return type:

tensor of shape (B, N, n_outputs)

class aidsorb.modules.TNet(embed_dim)

Bases: Module

Spatial transformer network (STN) from the [PointNet] paper for performing the input and feature transform.

T-Net takes as input a (possibly embedded) point cloud of shape (dim, N) and regresses a (dim, dim) matrix. Each point in the point cloud has shape (dim,).

The input must be batched, i.e. have shape of (B, dim, N), where B is the batch size and N is the number of points in each point cloud.

Parameters:

embed_dim (int) – Embedding dimension.

Examples

>>> tnet = TNet(embed_dim=64)
>>> x = torch.randn((128, 64, 42))  # Shape (B, embed_dim, N).
>>> tnet(x).shape
torch.Size([128, 64, 64])

forward(x)

Return the regressed matrices.

Parameters:

x (tensor of shape (B, embed_dim, N))

Returns:

out – Regressed matrices.

Return type:

tensor of shape (B, embed_dim, embed_dim)

aidsorb.modules.conv1d_block(in_channels, out_channels, config_activation=None, **kwargs)

Return a 1D convolutional block.

The block has the following form:

block = nn.Sequential(
    conv_layer,
    nn.BatchNorm1d(out_channels),
    activation_fn
    )

Parameters:

• in_channels (int or None) – If None, the conv_layer is lazy initialized.

• out_channels (int)

• config_activation (dict, default=None) –

  Dictionary for configuring activation function. If None, the ReLU activation is used.

  • 'name' activation’s class name str

  • 'hparams' activation’s hyperparameters dict

• **kwargs – Valid keyword arguments for Conv1d.

Returns:

block

Return type:

torch.nn.Sequential

Examples

>>> inp, out = 4, 128
>>> x = torch.randn(32, 4, 100)  # Shape (B, in_channels, N).
>>> config_afn = {'name': 'LeakyReLU', 'hparams': {'negative_slope': 0.5}}

>>> # Default activation function (ReLU).
>>> block = conv1d_block(inp, out, kernel_size=1)
>>> block(x).shape
torch.Size([32, 128, 100])
>>> block[2]
ReLU()

>>> # Custom activation function.
>>> block = conv1d_block(inp, out, config_afn, kernel_size=1)
>>> block(x).shape
torch.Size([32, 128, 100])
>>> block[2]
LeakyReLU(negative_slope=0.5)

>>> # Lazy initialized.
>>> block = conv1d_block(None, out, kernel_size=1)
>>> block(x).shape
torch.Size([32, 128, 100])

aidsorb.modules.dense_block(in_features, out_features, config_activation=None, **kwargs)

Return a dense block.

The block has the following form:

block = nn.Sequential(
    linear_layer,
    nn.BatchNorm1d(out_features),
    activation_fn,
    )

Parameters:

• in_features (int or None) – If None, the linear_layer is lazy initialized.

• out_features (int)

• config_activation (dict, default=None) –

  Dictionary for configuring activation function. If None, the ReLU activation is used.

  • 'name' activation’s class name str

  • 'hparams' activation’s hyperparameters dict

• **kwargs – Valid keyword arguments for Linear.

Returns:

block

Return type:

torch.nn.Sequential

Examples

>>> inp, out = 3, 10
>>> x = torch.randn(64, inp)  # Shape (B, in_features).
>>> config_afn = {'name': 'SELU', 'hparams': {}}

>>> # Default activation function (ReLU).
>>> block = dense_block(inp, out)
>>> block(x).shape
torch.Size([64, 10])
>>> block[2]
ReLU()

>>> # Custom activation function.
>>> block = dense_block(inp, out, config_afn)
>>> block(x).shape
torch.Size([64, 10])
>>> block[2]
SELU()

>>> # Lazy initialized.
>>> block = dense_block(None, 16)
>>> block(x).shape
torch.Size([64, 16])