aidsorb.modules

This module provides torch.nn.Module’s for building the architectures in aidsorb.models.

Currently, the module provides the basic blocks for building the architecture from the [PointNet] paper.

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).

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

Bases: Module

Backbone of the PointNet model.

This block is responsible for obtaining the local and global features, which can then be passed to a task head for predictions. This block 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:

• 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.

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

Examples

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

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

forward(x)

Return the features and critical indices.

The type of the features is determined by local_feats.

Parameters:

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

Returns:

out –

• out[0] == features

• out[1] == critical_indices

Return type:

tuple of length 2

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

Bases: Module

Classification head from the [PointNet] paper.

Note

This head can be used for classification or regression.

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.

Note

This head can be used for segmentation.

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) – The 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 – The regressed matrices.

Return type:

tensor of shape (B, embed_dim, embed_dim)

aidsorb.modules.conv1d_block(in_channels, out_channels, **kwargs)

Return a 1D convolutional block.

The block has the following form:

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

Parameters:

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

• out_channels (int)

• **kwargs – Valid keyword arguments for torch.nn.Conv1d.

Returns:

block

Return type:

torch.nn.Sequential

See also

torch.nn.Conv1d

For a description of the parameters.

Examples

>>> x = torch.randn(32, 4, 100)  # Shape (B, C_in, N).
>>> block = conv1d_block(4, 128, kernel_size=1)
>>> block(x).shape  # Shape (B, C_out, N).
torch.Size([32, 128, 100])

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

aidsorb.modules.dense_block(in_features, out_features, **kwargs)

Return a dense block.

The block has the following form:

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

Parameters:

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

• out_features (int)

• **kwargs – Valid keyword arguments for torch.nn.Linear.

Returns:

block

Return type:

torch.nn.Sequential

See also

torch.nn.Linear

For a description of the parameters.

Examples

>>> x = torch.randn(64, 3)  # Shape (B, in_features).
>>> block = dense_block(3, 10)
>>> block(x).shape  # Shape (B, out_features).
torch.Size([64, 10])

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