tnet

T-Net alignment modules predicting affine transforms for point and feature spaces.

Classes:

TNet –

Transformation Network (T-Net) module as described in PointNet paper
DynamicTNet –

Dynamic graph-based Transformation Network as used in the DGCNN part segmentation model.

TNet ¶

TNet(
    local_channels: Union[int, Sequence[int]],
    global_channels: Union[int, Sequence[int]],
    k: int,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    dropout: float = 0.0,
    aggr: AggrType = "max",
)

Bases: Module

Transformation Network (T-Net) module as described in PointNet paper PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation.

T-Net predicts an affine transformation matrix that helps align input point clouds or feature spaces to a canonical space. This network acts as a mini-PointNet that takes points/features as input and outputs a transformation matrix.

This layer will apply the following transformation to the input:

\[ x' = x \cdot T \]

where \(T\) is the transformation matrix.

There are two variants of the T-Net in PointNet: 1. Spatial transform network (ST-Net): Operates on point coordinates (k=3) 2. Feature transform network (FT-Net): Operates on point features (k=64 typically)

Note

The transformation matrix is initialized as an identity matrix and adds a residual connection to help with optimization stability.

Parameters:

local_channels (Union[int, Sequence[int]]) –

Channels of the first MLP, before pooling.
global_channels (Union[int, Sequence[int]]) –

Channels of the second MLP, after pooling.
k (int) –

Dimension of input features to transform.
act (Union[str, Callable, None], default: 'relu' ) –

Activation function to use.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the activation function.
act_first (bool, default: False ) –

Whether to apply the activation function before the normalization.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization to use.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the normalization.
bias (bool, default: True ) –

Whether to use bias in the linear layers.
dropout (float, default: 0.0 ) –

Dropout rate.
aggr (AggrType, default: 'max' ) –

Aggregation method to use.

Methods:

forward –

Forward pass of the T-Net.

forward ¶

forward(x: Tensor, batch: Tensor) -> Tensor

Forward pass of the T-Net.

Predicts one \(k \times k\) transformation matrix per sample and applies it to the points.

Parameters:

x (Tensor) –

Input tensor of shape \((N, k)\) where \(N\) is the number of points and \(k\) is the dimension of the input features.
batch (Tensor) –

Batch indices of shape \((N,)\).

Returns:

Tensor –

Transformed features of shape \((N, k)\).

DynamicTNet ¶

DynamicTNet(
    edge_channels: Sequence[int],
    local_channels: Sequence[int],
    global_channels: Sequence[int],
    k: int,
    num_neighbors: int = 20,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    dropout: float = 0.0,
    aggr: AggrType = "max",
)

Bases: Module

Dynamic graph-based Transformation Network as used in the DGCNN part segmentation model.

Unlike TNet which applies a point-wise MLP, this variant first builds a kNN graph and processes edge features with a DynamicEdgeConv, matching the Transform_Net from antao97/dgcnn.pytorch.

Architecture

edge_conv: DynamicEdgeConv over kNN graph features [2*k, ...edge_channels]
local_nn: Point-wise MLP [edge_channels[-1], ...local_channels], then scatter max
global_nn: Global MLP [local_channels[-1], ...global_channels]
transform: Linear projection to k * k matrix (initialized as identity)

Parameters:

edge_channels (Sequence[int]) –

Hidden channels for the EdgeConv MLP (excluding the 2*k input).
local_channels (Sequence[int]) –

Channels for the point-wise MLP applied after the EdgeConv.
global_channels (Sequence[int]) –

Channels for the global MLP applied after pooling.
k (int) –

Dimension of input features to transform.
num_neighbors (int, default: 20 ) –

Number of nearest neighbors for the dynamic graph.
act (Union[str, Callable, None], default: 'relu' ) –

Activation function.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the activation function.
act_first (bool, default: False ) –

Whether to apply the activation before normalization.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization layer type.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the normalization layer.
bias (bool, default: True ) –

Whether to use bias in linear layers.
dropout (float, default: 0.0 ) –

Dropout rate.
aggr (AggrType, default: 'max' ) –

Aggregation method to use.

Methods:

forward –

Args:

forward ¶

forward(x: Tensor, batch: Tensor) -> Tensor

Parameters:

x (Tensor) –

Point features of shape \((N, k)\).
batch (Tensor) –

Batch indices of shape \((N,)\).

Returns:

Tensor –

Transformed features of shape \((N, k)\).