pointnet2_blocks

PointNet++ set abstraction and feature propagation blocks.

Classes:

SAModule –

Single-resolution set-abstraction block (PointNet++ SSG/MSG).
GlobalSAModule –

Global set-abstraction block: a shared MLP followed by a pool over each batch element.
FPModule –

Feature-propagation block (PointNet++): \(k\)-NN interpolation, skip concatenation, and an MLP.
PointNet2Conv –

PointNet++ grouping convolution on top of PyG's MessagePassing.
PointNet2SetAbstraction –

Set-abstraction block built from one PointNet2Conv per grouping scale.
PointNet2GlobalSetAbstraction –

Global set-abstraction block: a shared MLP followed by a pool over each batch element.
PointNet2FeaturePropagation –

K-NN interpolation + skip concatenation + MLP, as in

Functions:

ensure_msg_list –

Utility function to ensure that items are converted in nested lists compatible
ensure_msg_list_size –

Validate the length of a sequence, then nest it for Multi-Scale Grouping (MSG) compatibility.

SAModule ¶

SAModule(
    in_channels: int,
    channels: Sequence[Union[int, Sequence[int]]],
    *,
    ratio: Optional[float] = None,
    num_points: Optional[int] = None,
    radii: Union[float, Sequence[float]],
    num_neighbors: Union[int, Sequence[int]],
    spatial_dim: int = 3,
    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 = False,
    use_pos: bool = True,
    normalize_pos: bool = True,
    pos_first: bool = False,
    pool: PoolLike = "max",
    sort_neighbors: bool = False,
)

Bases: Module

Single-resolution set-abstraction block (PointNet++ SSG/MSG).

Note

SAModule / GlobalSAModule / FPModule form the canonical PointNet++ stack used by the registered models. The PointNet2* classes in this module are an alternative implementation of the same blocks on top of PyG's MessagePassing.

Parameters:

ratio (Optional[float], default: None ) –

Fractional farthest-point-sampling rate. Mutually exclusive with num_points.
num_points (Optional[int], default: None ) –

Absolute number of centroids to sample (e.g. VoteNet's fixed \(2048, 1024, \ldots\)). Exactly one of ratio / num_points must be given. A sample with fewer than num_points points yields repeated centroids (FPS samples with replacement to keep shapes stable).
pos_first (bool, default: False ) –

Concatenate the relative position before the grouped features (cat([rel_pos, x])) instead of after. VoteNet and the reference PointNet++ kernels use this order; keeping it a flag lets weights convert as a pure rename without a column swap.

GlobalSAModule ¶

GlobalSAModule(
    in_channels: int,
    channels: Sequence[int],
    spatial_dim: int = 3,
    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 = False,
    use_pos: bool = False,
    pos_first: bool = False,
    pool: PoolLike = "max",
)

Bases: Module

Global set-abstraction block: a shared MLP followed by a pool over each batch element.

Parameters:

use_pos (bool, default: False ) –

Concatenate the absolute point positions to x before the MLP. Unlike SAModule there is no sampled centroid to offset against, so the coordinates enter unnormalized (the reference PointNet++ GroupAll).
pos_first (bool, default: False ) –

Concatenate the positions before the features (cat([pos, x])) instead of after.

FPModule ¶

FPModule(
    in_channels: int,
    channels: Sequence[int],
    k: int = 3,
    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 = False,
    weighting: Literal["squared", "inverse"] = "squared",
    eps: float = 1e-16,
)

Bases: Module

Feature-propagation block (PointNet++): \(k\)-NN interpolation, skip concatenation, and an MLP.

Parameters:

in_channels (int) –

Number of input channels after the skip concatenation.
channels (Sequence[int]) –

Per-layer channel sizes of the MLP.
k (int, default: 3 ) –

Number of neighbors for the \(k\)-NN interpolation. PointNet++ uses \(k = 3\); \(k = 1\) copies the nearest coarse feature (RandLA-Net).
weighting (Literal['squared', 'inverse'], default: 'squared' ) –

Inverse-distance weighting scheme passed to knn_interpolate. Irrelevant when \(k = 1\).
eps (float, default: 1e-16 ) –

Numerical stability term added to the interpolation distances.

PointNet2Conv ¶

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

Bases: MessagePassing

PointNet++ grouping convolution on top of PyG's MessagePassing.

Each message concatenates the neighbor features with the relative position (cat([x_j, pos_j - pos_i])) and applies local_nn; messages are aggregated per centroid.

Parameters:

local_nn (Module) –

Network applied to each message of shape \((E, C + D)\).
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").

PointNet2SetAbstraction ¶

PointNet2SetAbstraction(
    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",
)

Bases: Module

Set-abstraction block built from one PointNet2Conv per grouping scale.

Farthest point sampling selects the centroids, a ball query gathers the neighbors of each centroid per scale, and the per-scale outputs are concatenated (Multi-Scale Grouping when channels is a nested sequence).

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

PointNet2GlobalSetAbstraction ¶

PointNet2GlobalSetAbstraction(
    in_channels: int,
    channels: 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,
    aggr: PoolLike = "max",
)

Bases: Module

Global set-abstraction block: a shared MLP followed by a pool over each batch element.

Parameters:

in_channels (int) –

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

Per-layer channel sizes of the MLP.
dropout (float, default: 0.0 ) –

Dropout rate inside the MLP.
aggr (PoolLike, default: 'max' ) –

Pooling operation applied per batch element.

PointNet2FeaturePropagation ¶

PointNet2FeaturePropagation(
    channels: Sequence[int],
    k: int = 3,
    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,
    plain_last: bool = True,
    weighting: Literal["squared", "inverse"] = "inverse",
)

Bases: Module

K-NN interpolation + skip concatenation + MLP, as in PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space.

The interpolated features are concatenated before the skip features (cat([interp, skip])). Models with the opposite upstream cat order (RandLA-Net's cat([skip, interp])) must swap the first linear layer's column blocks at conversion time to stay weight-compatible.

Parameters:

channels (Sequence[int]) –

Per-layer channel sizes for the post-concat MLP.
k (int, default: 3 ) –

Number of neighbors for the K-NN interpolation. PointNet++ uses \(k = 3\) with inverse-distance weighting; RandLA-Net uses \(k = 1\) (nearest only).
weighting (Literal['squared', 'inverse'], default: 'inverse' ) –

Inverse-distance weighting scheme passed to knn_interpolate. Irrelevant when \(k = 1\).

ensure_msg_list ¶

ensure_msg_list(
    items: Sequence[Any], extra_msg: str = ""
) -> List[List[List[Any]]]

Utility function to ensure that items are converted in nested lists compatible with Multi-Scale Grouping (MSG) mode. This function will convert a list of list into a list of list of list.

Example

Let's say we have designed a network where the first two SA blocks are not using MSG mode, but the last SA block is using MSG mode.

Calling ensure_msg_list will make sure the provided channels are compliant with the MSG mode.

>>> sa_channels = [[32, 64], [128, 256], [[256, 512, 512], [256, 512, 1024]]]
>>> ensure_msg_list(sa_channels)
[[[32, 64]], [[128, 256]], [[256, 512, 512], [256, 512, 1024]]]

ensure_msg_list_size ¶

ensure_msg_list_size(
    value: Sequence[Any], size: int, extra_msg: str = ""
) -> Sequence[Any]

Validate the length of a sequence, then nest it for Multi-Scale Grouping (MSG) compatibility.

Parameters:

value (Sequence[Any]) –

Sequence of per-block channel specifications.
size (int) –

Expected number of elements in value.
extra_msg (str, default: '' ) –

Extra context appended to the error message.

Returns:

Sequence[Any] –

The value converted to a list of lists of lists (one inner list per grouping scale).

Raises:

ValueError –

If value does not have exactly size elements.

Example

>>> ensure_msg_list_size([[32, 64], [64, 128]], size=2)
[[[32, 64]], [[64, 128]]]