pointnext_blocks

PointNeXt convolution layer introduced in the PointNeXt: Revisiting PointNet++ with Improved Training and Scaling Strategies by Guocheng Qian et al.

Note

This layer is also referred to as Local Aggregation in different papers and implementations.

This layer implements the torch_geometric.nn.conv.MessagePassing interface from PyTorch Geometric, which allows for local aggregation of features.

Tip

This layer is similar to the torch_geometric.nn.conv.PointNetConv layer from PyTorch Geometric, and introduces relative position normalization.

You can use it as follows:

import torch
from torch_geometric.nn import MLP, radius_graph
from torch_pointcloud.layers import PointNeXtConv

torch.manual_seed(0)
x = torch.randn(10, 10)
pos = torch.randn(10, 3)
batch = torch.zeros(10, dtype=torch.long)
edge_index = radius_graph(pos, r=1.5, batch=batch, max_num_neighbors=16)

conv = PointNeXtConv(MLP([3 + 10, 10]))

# Normalize the relative position by the query radius
out = conv(x, pos, edge_index, pos_divisor=1.5)

# This will be equivalent to the PointNetConv layer
out = conv(x, pos, edge_index)

Classes:

PointNeXtConv –

PointNeXt grouping convolution on top of PyG's MessagePassing.
PointNeXtSetAbstraction –

PointNeXt set-abstraction block: FPS centroids, radius-normalized grouping, and optional
PointNeXtResidualBlock –

PointNeXt inverted-residual MLP block (InvResMLP).

PointNeXtConv ¶

PointNeXtConv(
    local_nn: Module,
    add_self_loops: bool = True,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

PointNeXt grouping convolution on top of PyG's MessagePassing.

Each message concatenates the relative position with the neighbor features (cat([pos_j - pos_i, x_j])) and applies local_nn; when a pos_divisor is given the relative positions are normalized by it (PointNeXt normalizes by the ball-query radius).

Parameters:

local_nn (Module) –

Network applied to each message of shape \((E, D + C)\).
add_self_loops (bool, default: True ) –

Whether to add self-loops to the edge index.
**kwargs (Unpack[MessagePassingParams], default: {} ) –

Additional MessagePassing arguments (aggr defaults to "max").

PointNeXtSetAbstraction ¶

PointNeXtSetAbstraction(
    spatial_dim: int,
    in_channels: int,
    channels: Sequence[Union[int, Sequence[int]]],
    ratio: float,
    radius: Union[float, Sequence[float]],
    num_neighbors: Union[int, Sequence[int]],
    dropout: float = 0.0,
    act: Union[str, Callable, None] = "relu",
    act_first: bool = False,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: Union[bool, List[bool]] = True,
    add_self_loops: bool = False,
    aggr: AggrType = "max",
    use_res: bool = True,
)

Bases: Module

PointNeXt set-abstraction block: FPS centroids, radius-normalized grouping, and optional residual skip connections.

Farthest point sampling selects the centroids, a ball query gathers the neighbors of each centroid per scale (Multi-Scale Grouping when channels is a nested sequence), and the relative positions are normalized by the query radius before the grouping convolution.

Parameters:

spatial_dim (int) –

Dimension of point coordinates.
in_channels (int) –

Number of input feature channels.
channels (Sequence[Union[int, Sequence[int]]]) –

Per-scale MLP channel sizes; a nested sequence enables Multi-Scale Grouping.
ratio (float) –

Fractional farthest-point-sampling rate.
radius (Union[float, Sequence[float]]) –

Ball-query radius per scale, also used to normalize the relative positions.
num_neighbors (Union[int, Sequence[int]]) –

Maximum number of neighbors per scale.
dropout (float, default: 0.0 ) –

Dropout rate inside the per-scale MLPs.
add_self_loops (bool, default: False ) –

Whether to add self-loops to the grouping edge index.
aggr (AggrType, default: 'max' ) –

Message aggregation used by the convolutions.
use_res (bool, default: True ) –

Whether each scale adds a residual skip connection from the sampled centroid features (with a linear projection when the channel counts differ).

PointNeXtResidualBlock ¶

PointNeXtResidualBlock(
    spatial_dim: int,
    channels: int,
    expansion: int,
    radius: float,
    num_neighbors: 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: Union[bool, List[bool]] = True,
    add_self_loops: bool = False,
    aggr: AggrType = "max",
)

Bases: Module

PointNeXt inverted-residual MLP block (InvResMLP).

A radius-graph grouping convolution followed by an inverted-bottleneck MLP (channels -> channels * expansion -> channels), wrapped in a single residual connection. The resolution is unchanged; downsampling happens in PointNeXtSetAbstraction.

Parameters:

spatial_dim (int) –

Dimension of point coordinates.
channels (int) –

Number of input and output feature channels.
expansion (int) –

Expansion factor of the bottleneck MLP.
radius (float) –

Radius of the grouping graph, also used to normalize the relative positions.
num_neighbors (int) –

Maximum number of neighbors in the grouping graph.
add_self_loops (bool, default: False ) –

Whether to add self-loops to the grouping edge index.
aggr (AggrType, default: 'max' ) –

Message aggregation used by the convolution.