Overview

Reusable neural network blocks shared across the model architectures.

Modules:

act –

Activation factory wrapper around torch_geometric.nn.resolver.
affine –

Per-channel affine transformation as a function and a learnable module.
anchors –

Anchor-based dense detection heads for the voxel detectors (PointPillars, SECOND).
bev_backbone –

SSD-style 2D BEV backbones shared by the voxel detectors (PointPillars, SECOND, Voxel Mamba).
conv2d_blocks –

2D convolution block with optional normalization and activation.
conv3d_blocks –

3D convolution block with optional normalization and activation.
dropouts –

Dropout layers: stochastic depth (DropPath) and the create_dropout factory.
fps –

Farthest point sampling (FPS) as a module returning sampled indices.
geometric_affine –

PointMLP grouping convolution with geometric affine neighborhood normalization.
grid_pool –

Grid-based point cloud pooling (downsampling).
grouped_vector_attention –

Grouped vector attention over a neighborhood graph.
kpconv_blocks –

Kernel point convolution, its kernel point placement, and the normalized and residual blocks built on it.
linear_blocks –

Linear block with optional normalization and activation.
norms –

Normalization factory wrapper with spatial dimensionality support.
octree_attention –

OctFormer octree window attention with relative position encoding.
octree_blocks –

Octree convolution and deconvolution blocks with normalization and activation.
pdnorm –

Prompt-driven normalization (PDNorm) routing each batch through a per-condition norm.
point_patch_embed –

Mini-PointNet patch embedding turning local point groups into tokens.
point_transformer_conv –

Point Transformer vector-attention convolution over a neighborhood graph.
pointconv_blocks –

PointConv message-passing convolutions and the set-abstraction blocks built on them.
pointnet2_blocks –

PointNet++ grouping convolution, set abstraction and feature propagation blocks.
pointnext_blocks –

PointNeXt convolution layer introduced in the
pools –

Per-segment pooling modules over packed batches and the create_pool / create_adaptive_pool factories.
pvcnn_blocks –

PVCNN building blocks: voxelization, squeeze-and-excitation, and point-voxel convolution.
randlanet_blocks –

RandLA-Net local feature aggregation: spatial encoding, attentive pooling, and the dilated residual block.
rope –

3D rotary position embedding for point-cloud attention.
serialized_attention –

Serialized attention variants used by Point Transformer V3 and descendants.
serialized_pool –

Pooling and upsampling driven by point cloud serialization codes.
spconv_blocks –

Sparse convolution blocks: submanifold, strided, and residual variants.
tnet –

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

torchsparse convolution blocks: a conv + norm + act block and a residual block over SparseTensors.
transformer –

Transformer building blocks for point-cloud backbones and detectors.
vfe –

Voxel feature encoders: the pillar feature-net layer and the dynamic mean voxel encoder.
view –

Tensor reshaping as a module wrapping Tensor.view.
voxel_grid_pool –

Voxel grid subsampling of packed point clouds.
xconv –

PointCNN XConv layer: convolution on X-transformed point neighborhoods.

Classes:

Affine –

Applies an affine transformation to the input.
AnchorGroupHead –

One RPN head of MultiGroupAnchorHead.
AnchorHead –

Single-stage anchor head.
AnchorHeadMultiOutput –

Predictions of MultiGroupAnchorHead.
AnchorHeadOutput –

Raw and decoded predictions of AnchorHead.
MultiGroupAnchorHead –

Multi-group anchor head, one separate head per group of classes.
BEVBackbone –

SSD-style multi-scale 2D BEV backbone.
BEVResidualBackbone –

Residual SSD-style 2D BEV backbone used by Voxel Mamba.
ResidualBlock2d –

Residual 2D conv block (the reference's BasicBlock) of the BEV residual backbone.
Conv2dBlock –

Single nn.Conv2d (or nn.ConvTranspose2d) + optional norm + optional activation.
Conv3dBlock –

Single nn.Conv3d + optional nn.BatchNorm3d + optional activation.
DropPath –

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
FPS –

Farthest point sampling as a module, returning the sampled indices.
GeometricAffineConv –

Grouping convolution with geometric affine normalization (PointMLP) from
GridPool –

Grid-based downsampling that clusters points by quantized coordinates.
GroupedVectorAttention –

Vector attention over a neighborhood graph, with one weight vector shared by each group of channels.
KPConv –

Kernel point convolution over a neighborhood graph.
KPConvBlock –

Kernel point convolution followed by normalization and activation.
KPResidualBlock –

Bottleneck residual block: a channel-reducing MLP, a KPConvBlock, and a channel-restoring MLP.
LinearBlock –

Linear block consisting of a linear layer, normalization and activation.
OctreeAttention –

Multi-head self-attention restricted to windows of patch_size consecutive octree nodes.
OctreeRelativePositionEncoding –

Relative Position Encoding (RPE) module used within the OctreeAttention module.
OctreeT –

An enhanced Octree with transformer-specific capabilities (patching, dilation, masking).
OctreeConvBlock –

Octree convolution followed by normalization and activation.
OctreeDeconvBlock –

Octree transposed convolution followed by normalization and activation.
PDNorm –

Prompt-driven normalization that routes each batch through a per-condition norm.
PointPatchEmbed –

Mini-PointNet patch (token) embedding for grouped point clouds.
PointTransformerConv –

The Point Transformer layer from the
PointConv –

PointConv message passing: a per-edge feature outer-producted with a learned continuous weight.
PointConvDensity –

PointConv with inverse density reweighting.
PointConvDensityGlobalSetAbstraction –

Global PointConv set-abstraction block with inverse-density re-weighting.
PointConvDensitySetAbstraction –

PointConv set-abstraction block with inverse-density re-weighting.
PointConvGlobalSetAbstraction –

Global PointConv set-abstraction block: one weight-net convolution over each whole sample.
PointConvSetAbstraction –

PointConv set-abstraction block: optional downsampling, \(k\)-NN grouping, and a
PointNet2Conv –

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

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

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

Set-abstraction block (PointNet++ SSG / MSG) built from one PointNet2Conv per grouping scale.
PointNeXtConv –

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

PointNeXt inverted-residual MLP block (InvResMLP).
PointNeXtSetAbstraction –

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

Runs several pools on the same input and concatenates their outputs along the feature dim.
MaxPool –

Per-segment max pooling over a packed batch, via torch_geometric.utils.scatter(reduce="max").
MeanPool –

Per-segment mean pooling over a packed batch, via torch_geometric.utils.scatter(reduce="mean").
MinPool –

Per-segment min pooling over a packed batch, via torch_geometric.utils.scatter(reduce="min").
MulPool –

Per-segment product pooling over a packed batch, via torch_geometric.utils.scatter(reduce="mul").
SumPool –

Per-segment sum pooling over a packed batch, via torch_geometric.utils.scatter(reduce="sum").
PVConv –

Point-voxel convolution: a dense 3D conv branch summed with a per-point MLP branch.
SE3d –

Squeeze-and-excitation gate for dense voxel grids.
Voxelization –

Averages packed point features into a dense voxel grid.
AttentivePooling –

Attention-weighted aggregation of neighbor features.
LocalFeatureAggregation –

Local feature aggregation: two rounds of spatial encoding and attentive pooling.
LocalSpatialEncoding –

Per-edge spatial encoding MLP.
RandLANetResidualBlock –

RandLA-Net dilated residual block.
Point3DRoPE –

3D Rotary Position Embedding for point cloud attention.
PatchLayout –

Padding and patch offsets of one point set for serialized attention.
RelativePositionalEncoding –

Table-based relative position bias added to the attention logits of a serialized patch.
SerializedAttention –

Vanilla serialized attention from
SerializedAttentionRoPE –

Serialized attention with 3D rotary position embedding from
SerializedAttentionRPE –

Serialized attention with the relative position bias from PT-V3.
SerializedPool –

Grid pooling driven by the serialization code: points sharing a coarser code are reduced to one.
SerializedUpsample –

Undoes a SerializedPool step: scatters the pooled features back and adds the projected skip.
SparseBasicBlock –

Residual block of two submanifold sparse convolutions: conv, norm, act, conv, norm, add the input, act.
SparseConvBlock –

Sparse 3D convolution followed by normalization and activation.
SparseResidualBlock –

Pre-activation sparse residual block shared by the SPFormer and SphereFormer SpConv U-Nets.
SubMConv3dBlock –

Submanifold sparse 3D convolution followed by normalization and activation, on packed point features.
SubMConv3dResidualBlock –

Residual block with a single \(3\times3\times3\) submanifold convolution: conv, norm, add the input, act.
DynamicTNet –

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

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

Sparse 3D convolution followed by normalization and activation.
TorchSparseResidualBlock –

Residual block of two sparse 3D convolutions.
Attention –

Multi-head self-attention over a dense token sequence.
TransformerBlock –

Pre-norm transformer block: residual multi-head attention then a residual MLP.
TransformerDecoderLayer –

DETR-style transformer decoder layer: self-attention over the queries, cross-attention onto the memory and a
TransformerEncoderLayer –

DETR-style transformer encoder layer: self-attention with positional embeddings added to the query and key,
DynamicMeanVFE –

Dynamic mean voxel feature encoder for the Voxel Mamba detector.
PillarFeatureLayer –

Pillar feature-net layer: linear + norm + act on every point, then a max-pool over each pillar.
View –

Returns a new view of the tensor with the same data.
VoxelGridPool –

Voxel grid subsampling: the points falling in the same grid_size cell are reduced to a single point.
XConv –

XConv layer as described in the paper

Functions:

create_module –

Utility function to instantiate modules from a registry, or directly pass a module instance.
create_act –

Resolve an activation, returning None for act=None.
prediction_branch –

Build a prediction branch: middle conv blocks, then a plain output conv.
create_dropout –

Resolve a dropout layer from a name, a class, an instance, or a probability, returning None for name=None.
drop_path –

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks),
create_kernel_points –

Builds the kernel point positions of a KPConv kernel, randomly rotated and jittered.
create_norm –

Resolve a normalization layer with a spatial dimensionality hint.
create_adaptive_pool –

Resolve a dense adaptive pooling module (nn.AdaptiveAvgPool1d / nn.AdaptiveMaxPool1d), returning None
create_pool –

Resolve a packed-batch pooling module from a name, class, or instance, returning None for name=None.
patch_layout –

Pad every sample to whole patches of patch_size points and derive the patch offsets.

Affine [source] ¶

Affine(
    num_features: int,
    bias: bool = True,
    device: Optional[device] = None,
    dtype: Optional[dtype] = None,
)

Bases: Module

Applies an affine transformation to the input. This layer will apply the following transformation to the input tensor \(x\):

\[ y = x \cdot \text{weight} + \text{bias} \]

Parameters:

num_features (int) –

The number of features in the input.
bias (bool, default: True ) –

Whether to use bias.
device (Optional[device], default: None ) –

The device to use.
dtype (Optional[dtype], default: None ) –

The dtype to use.

Shape

Input: \((N, *, C)\) where \(*\) means any number of additional dimensions.
Output: \((N, *, C)\) where \(*\) means any number of additional dimensions.

AnchorGroupHead [source] ¶

AnchorGroupHead(
    input_channels: int,
    num_classes: int,
    num_anchors_per_location: int,
    code_size: int,
    reg_list: Sequence[str],
    head_label_indices: Tensor,
    *,
    num_middle_conv: int = 1,
    num_middle_filter: int = 64,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

One RPN head of MultiGroupAnchorHead.

A SeparateHead-style head over the shared feature: a classification branch plus one regression branch per box-code group (reg, height, size, angle, velo), whose outputs are concatenated into the per-anchor box code. Predictions are reshaped to the multihead layout \((B, A, \cdot)\) with \(A\) the anchors of this head's class group.

Parameters:

input_channels (int) –

Channels of the shared feature map.
num_classes (int) –

Number of classes handled by this head (separate-multihead).
num_anchors_per_location (int) –

Anchors per BEV cell for this head.
code_size (int) –

Box code size (e.g. 10 for sincos angle + velocity).
reg_list (Sequence[str]) –

Regression-branch spec, e.g. ["reg:2", "height:1", "size:3", "angle:2", "velo:2"].
num_middle_conv (int, default: 1 ) –

Number of middle convs per branch.
num_middle_filter (int, default: 64 ) –

Middle-conv channel width.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the middle convs.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable for the middle convs.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

Methods:

reset_parameters –

Start the class logits at the focal-loss prior of \(1\%\) foreground anchors.

reset_parameters [source] ¶

reset_parameters() -> None

Start the class logits at the focal-loss prior of \(1\%\) foreground anchors.

AnchorHead [source] ¶

AnchorHead(
    input_channels: int,
    num_classes: int,
    grid_size: Tuple[int, int],
    point_cloud_range: Sequence[float],
    *,
    anchor_sizes: Sequence[Sequence[float]],
    anchor_bottom_heights: Sequence[float],
    feature_map_stride: int,
    anchor_rotations: Sequence[float] = (0.0, 1.57),
    num_dir_bins: int = 2,
    dir_offset: float = 0.78539,
    dir_limit_offset: float = 0.0,
)

Bases: Module

Single-stage anchor head.

Three \(1\times1\) convs predict, per anchor, class logits, 7-DoF box residuals and a direction bin. At inference the residuals are decoded against the precomputed anchors and the predicted heading is snapped to the predicted direction bin (dir_offset / num_dir_bins).

Parameters:

input_channels (int) –

Channels of the BEV feature map fed to the head.
num_classes (int) –

Number of foreground classes.
grid_size (Tuple[int, int]) –

Full voxel grid size \((n_x, n_y)\) (before the head feature-map stride).
point_cloud_range (Sequence[float]) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
anchor_sizes (Sequence[Sequence[float]]) –

Per-class box size \((d_x, d_y, d_z)\), shape \((\text{num\_classes}, 3)\).
anchor_bottom_heights (Sequence[float]) –

Per-class anchor bottom \(z\), shape \((\text{num\_classes},)\).
anchor_rotations (Sequence[float], default: (0.0, 1.57) ) –

Yaw angles (radians) shared by all classes.
feature_map_stride (int) –

BEV feature-map stride of the head.
num_dir_bins (int, default: 2 ) –

Number of direction bins.
dir_offset (float, default: 0.78539 ) –

Direction-classifier angle offset.
dir_limit_offset (float, default: 0.0 ) –

Offset used when wrapping the decoded heading before snapping.

Methods:

reset_parameters –

Start the class logits at the focal-loss prior of \(1\%\) foreground anchors and the box residuals near zero.
generate_predicted_boxes –

Decode raw head outputs into per-anchor class logits and absolute boxes.
decode –

Decode a forward output into raw per-anchor detections (no score threshold or NMS).

reset_parameters [source] ¶

reset_parameters() -> None

Start the class logits at the focal-loss prior of \(1\%\) foreground anchors and the box residuals near zero.

generate_predicted_boxes [source] ¶

generate_predicted_boxes(
    batch_size: int,
    cls_preds: Tensor,
    box_preds: Tensor,
    dir_cls_preds: Tensor,
) -> Tuple[Tensor, Tensor]

Decode raw head outputs into per-anchor class logits and absolute boxes.

Returns:

Tensor –

A tuple (batch_cls_preds, batch_box_preds) of shapes \((B, A, \text{num\_classes})\)
Tensor –

(raw logits, not sigmoided) and \((B, A, 7)\) where \(A\) is the number of anchors.

decode [source] ¶

decode(out: AnchorHeadOutput) -> Detection3D

Decode a forward output into raw per-anchor detections (no score threshold or NMS).

Scores each anchor by its top sigmoid class probability and labels it by the argmax class. The full per-anchor set is returned; the evaluation pipeline applies score thresholding and per-class 3D NMS via the torch_pointcloud.ops.box3d utilities (see the benchmark examples).

Returns:

Detection3D –

Packed per-anchor detections `{"boxes": (B * A, 7), "scores": (B * A,), "labels": (B * A,),
Detection3D –

"batch": (B * A,)}` (PyG layout).

AnchorHeadMultiOutput [source] ¶

Bases: TypedDict

Predictions of MultiGroupAnchorHead.

AnchorHeadOutput [source] ¶

Bases: TypedDict

Raw and decoded predictions of AnchorHead.

MultiGroupAnchorHead [source] ¶

MultiGroupAnchorHead(
    input_channels: int,
    num_classes: int,
    grid_size: Tuple[int, int],
    point_cloud_range: Sequence[float],
    *,
    anchor_sizes: Sequence[Sequence[float]],
    anchor_bottom_heights: Sequence[float],
    head_class_groups: Sequence[Sequence[int]],
    feature_map_stride: int,
    anchor_rotations: Sequence[float] = (0.0, 1.57),
    shared_conv_num_filter: int = 64,
    reg_list: Sequence[str] = (
        "reg:2",
        "height:1",
        "size:3",
        "angle:2",
        "velo:2",
    ),
    num_middle_conv: int = 1,
    num_middle_filter: int = 64,
    code_size: int = 9,
    encode_angle_by_sincos: bool = True,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Multi-group anchor head, one separate head per group of classes.

A shared conv feeds several AnchorGroupHeads, one per class group. Anchors (7-DoF, padded to the box-code size) are decoded with decode_box_residuals (sincos heading, velocity deltas); per-head class scores stay separate (with their global label mapping) for class-wise NMS downstream.

Parameters:

input_channels (int) –

Channels of the BEV feature map fed to the head.
num_classes (int) –

Number of foreground classes.
grid_size (Tuple[int, int]) –

Full voxel grid size \((n_x, n_y)\).
point_cloud_range (Sequence[float]) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
anchor_sizes (Sequence[Sequence[float]]) –

Per-class box size \((d_x, d_y, d_z)\), shape \((\text{num\_classes}, 3)\).
anchor_bottom_heights (Sequence[float]) –

Per-class anchor bottom \(z\), shape \((\text{num\_classes},)\).
head_class_groups (Sequence[Sequence[int]]) –

Class-index groups, one per RPN head (e.g. [[0], [1, 2], ...]); the classes in each group share one SeparateHead.
anchor_rotations (Sequence[float], default: (0.0, 1.57) ) –

Yaw angles (radians) shared by all classes.
feature_map_stride (int) –

BEV feature-map stride of the head.
shared_conv_num_filter (int, default: 64 ) –

Channels of the shared conv.
reg_list (Sequence[str], default: ('reg:2', 'height:1', 'size:3', 'angle:2', 'velo:2') ) –

Regression-branch spec for each head.
num_middle_conv (int, default: 1 ) –

Middle convs per branch.
num_middle_filter (int, default: 64 ) –

Middle-conv channel width.
code_size (int, default: 9 ) –

Base box code size (9 for nuScenes; +1 internally for sincos).
encode_angle_by_sincos (bool, default: True ) –

Encode heading as \((\cos, \sin)\).
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the shared conv and head middle convs.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable for the shared conv and head middle convs.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

Methods:

decode –

Decode a multihead forward output into raw per-anchor detections (no score threshold or NMS).

decode [source] ¶

decode(out: AnchorHeadMultiOutput) -> Detection3D

Decode a multihead forward output into raw per-anchor detections (no score threshold or NMS).

Each head scores its anchors by their top sigmoid class probability and maps the argmax to the global label; the per-head results are concatenated in head order (matching batch_box's anchor order). When the box code carries velocity deltas the decoded \((v_x, v_y)\) columns are returned under velocity. The full per-anchor set is returned; the evaluation pipeline applies score thresholding and per-class 3D NMS via the torch_pointcloud.ops.box3d utilities (see the benchmark examples).

Returns:

Detection3D –

Packed per-anchor detections `{"boxes": (B * A, 7), "scores": (B * A,), "labels": (B * A,),
Detection3D –

"batch": (B * A,)}(PyG layout), plus"velocity"` \((B \cdot A, 2)\) when the head predicts it.

BEVBackbone [source] ¶

BEVBackbone(
    input_channels: int,
    layer_nums: Sequence[int],
    layer_strides: Sequence[int],
    num_filters: Sequence[int],
    upsample_strides: Sequence[float],
    num_upsample_filters: Sequence[int],
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

SSD-style multi-scale 2D BEV backbone.

Each level downsamples the BEV pseudo-image with a strided \(3\times3\) conv followed by layer_nums residual-free \(3\times3\) convs, then upsamples back to a common stride; the level outputs are concatenated along the channel dim. An upsample factor \(\geq 1\) uses a transposed conv, a factor \(< 1\) (e.g. \(0.5\)) a strided down-conv (nuScenes configs use both).

Parameters:

input_channels (int) –

Channels of the input BEV feature map.
layer_nums (Sequence[int]) –

Number of \(3\times3\) convs after the strided conv, per level.
layer_strides (Sequence[int]) –

Downsample stride of the leading conv, per level.
num_filters (Sequence[int]) –

Channel width, per level.
upsample_strides (Sequence[float]) –

Upsample factor per level.
num_upsample_filters (Sequence[int]) –

Channels of each upsampled level.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for every conv block.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable for every conv block.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

BEVResidualBackbone [source] ¶

BEVResidualBackbone(
    input_channels: int,
    layer_nums: Sequence[int],
    layer_strides: Sequence[int],
    num_filters: Sequence[int],
    upsample_strides: Sequence[float],
    num_upsample_filters: Sequence[int],
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Residual SSD-style 2D BEV backbone used by Voxel Mamba.

Same scaffolding as BEVBackbone (per-level block then upsample, concatenated), but each level is a stack of residual ResidualBlock2ds instead of plain \(3\times3\) convs.

Parameters:

input_channels (int) –

Channels of the input BEV feature map.
layer_nums (Sequence[int]) –

Number of residual blocks after the strided block, per level.
layer_strides (Sequence[int]) –

Downsample stride of the leading block, per level.
num_filters (Sequence[int]) –

Channel width, per level.
upsample_strides (Sequence[float]) –

Upsample factor per level.
num_upsample_filters (Sequence[int]) –

Channels of each upsampled level.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for every conv block.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable for every conv block.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

ResidualBlock2d [source] ¶

ResidualBlock2d(
    in_channels: int,
    out_channels: int,
    *,
    stride: int = 1,
    downsample: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Residual 2D conv block (the reference's BasicBlock) of the BEV residual backbone.

Parameters:

in_channels (int) –

Input channels.
out_channels (int) –

Output channels.
stride (int, default: 1 ) –

Stride of the first conv (and the optional projection shortcut).
downsample (bool, default: False ) –

Add a \(1\times1\) projection shortcut to match channels / stride.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

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

Extra normalization arguments.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra activation arguments.

Conv2dBlock [source] ¶

Conv2dBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    *,
    stride: int = 1,
    padding: int = 0,
    transposed: bool = False,
    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: bool = False,
)

Bases: Module

Single nn.Conv2d (or nn.ConvTranspose2d) + optional norm + optional activation.

The 2D analogue of Conv3dBlock, with stride, padding and transposed exposed so it can express the strided down-convs and transposed up-convs of an SSD-style BEV backbone.

Shape

Input: \((B, C_\text{in}, H, W)\) Output: \((B, C_\text{out}, H', W')\)

Parameters:

in_channels (int) –

Input channel count.
out_channels (int) –

Output channel count.
kernel_size (int, default: 3 ) –

Conv kernel size.
stride (int, default: 1 ) –

Conv stride.
padding (int, default: 0 ) –

Conv padding.
transposed (bool, default: False ) –

Use nn.ConvTranspose2d instead of nn.Conv2d (for upsampling).
act (Union[str, Callable, None], default: 'relu' ) –

Activation, name resolved by create_act. None disables.
act_first (bool, default: False ) –

If True, run activation before normalization.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the activation.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization, name resolved by create_norm (with dim=2). None disables.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the normalization.
bias (bool, default: False ) –

Whether the conv has a bias term.

Conv3dBlock [source] ¶

Conv3dBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    *,
    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: bool = True,
)

Bases: Module

Single nn.Conv3d + optional nn.BatchNorm3d + optional activation.

Shape

Input: \((B, C_\text{in}, R, R, R)\) Output: \((B, C_\text{out}, R, R, R)\)

Parameters:

in_channels (int) –

Input channel count.
out_channels (int) –

Output channel count.
kernel_size (int, default: 3 ) –

Conv3d kernel size.
act (Union[str, Callable, None], default: 'relu' ) –

Activation, name resolved by create_act. None disables.
act_first (bool, default: False ) –

If True, run activation before normalization.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the activation.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization, name resolved by create_norm. None disables.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the normalization.
bias (bool, default: True ) –

Whether the Conv3d has a bias term.

DropPath [source] ¶

DropPath(
    drop_prob: float = 0.0, scale_by_keep: bool = True
)

Bases: Module

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

FPS [source] ¶

FPS(ratio: float, random_start: Optional[bool] = None)

Bases: Module

Farthest point sampling as a module, returning the sampled indices.

Parameters:

ratio (float) –

Fraction of points to keep per sample.
random_start (Optional[bool], default: None ) –

Whether farthest point sampling starts from a random point. With the default None, the start is random in training mode and pinned in eval mode, so eval predictions are reproducible across runs.

Shape

Input: positions of shape \((N, D)\) with an optional batch vector of shape \((N,)\).
Output: sampled indices of shape \((M,)\) with \(M \approx N \cdot \text{ratio}\).

GeometricAffineConv [source] ¶

GeometricAffineConv(
    local_nn: Module,
    channels: int,
    spatial_dim: int = 3,
    use_pos: bool = True,
    normalize: NormalizeType = "center",
    std_mode: StdModeType = "graph",
    add_self_loops: bool = False,
    eps: float = 1e-05,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

Grouping convolution with geometric affine normalization (PointMLP) from Rethinking Network Design and Local Geometry in Point Cloud: A Simple Residual MLP Framework by Xu Ma, Can Qin, Haoxuan You, Haoxi Ran, Yun Fu.

Each neighborhood is recentered (normalize="center" subtracts the neighborhood mean, "anchor" subtracts the center's values), rescaled by the standard deviation of the offsets over the graph (std_mode="graph") or over the whole batch (std_mode="batch"), and modulated by learnable affine parameters; the center features are appended to each normalized message before local_nn and aggregation. The batch scope reproduces the normalization of the original PointMLP release, whose checkpoints expect it; it makes the output of one sample depend on the other samples in the batch.

Parameters:

local_nn (Module) –

Network applied to each message.
channels (int) –

Number of input feature channels.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.
use_pos (bool, default: True ) –

Whether to concatenate the point positions to the features before normalization.
normalize (NormalizeType, default: 'center' ) –

Neighborhood normalization mode ("center" or "anchor").
std_mode (StdModeType, default: 'graph' ) –

Scope of the standard deviation that rescales the offsets: one per graph ("graph") or one over every graph in the batch ("batch").
add_self_loops (bool, default: False ) –

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

Additional MessagePassing arguments.

GridPool [source] ¶

GridPool(
    in_channels: int,
    out_channels: int,
    stride: int = 2,
    bias: bool = True,
    act: Union[str, Callable, None] = None,
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    reduce: str = "max",
)

Bases: Module

Grid-based downsampling that clusters points by quantized coordinates.

Divides pos_grid by stride, groups unique voxels via torch.unique, and reduces features with torch_geometric.utils.segment.

Parameters:

in_channels (int) –

Number of input feature channels.
out_channels (int) –

Number of output feature channels.
stride (int, default: 2 ) –

Spatial stride for grid quantization.
act (Union[str, Callable, None], default: None ) –

Activation layer applied after projection.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

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

Apply activation before normalization.
norm (Union[str, Callable, None], default: None ) –

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

Extra arguments for the normalization layer.
reduce (str, default: 'max' ) –

Scatter reduction ("max", "mean", "sum", "min").

Methods:

forward –

Downsample points by grid quantization.

forward [source] ¶

forward(
    x: Tensor,
    pos_grid: Tensor,
    batch: Tensor,
    pos: Optional[Tensor] = None,
    condition: Optional[str] = None,
) -> Tuple[
    Tensor, Tensor, Tensor, Tensor, Optional[Tensor]
]

Downsample points by grid quantization.

Parameters:

x (Tensor) –

Point features of shape \((N, C_{in})\).
pos_grid (Tensor) –

Integer grid coordinates of shape \((N, 3)\). Must be non-negative and pool to coordinates below \(2^{16}\): batch index and coordinates are bit-packed into a single int64 cluster key, so out-of-range coordinates would merge clusters across samples.
batch (Tensor) –

Batch indices of shape \((N,)\).
pos (Optional[Tensor], default: None ) –

Optional real-valued positions of shape \((N, 3)\) to mean-pool alongside the features (e.g. for downstream rotary position embedding).
condition (Optional[str], default: None ) –

Optional condition name selecting the inner norm when the norm is a PDNorm.

Returns:

Tensor –

Tuple of (x_pooled, pos_grid_pooled, batch_pooled, pooling_inverse, pos_pooled).
Tensor –

pooling_inverse maps each input point to its pooled cluster index.
Tensor –

pos_pooled is None when pos is not provided.

GroupedVectorAttention [source] ¶

GroupedVectorAttention(
    channels: int,
    num_groups: int,
    attn_drop: float = 0.0,
    qkv_bias: bool = True,
    pe_multiplier: bool = False,
    pe_bias: bool = True,
    norm: Union[str, Callable, None] = "batch_norm",
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

Vector attention over a neighborhood graph, with one weight vector shared by each group of channels.

The relation between a query and its neighbor keys is optionally scaled and shifted by a learned encoding of their relative position, then mapped to num_groups weights and softmax-normalized over each destination's neighbors.

edge_index lists [source, target] pairs: keys and values come from the source points and queries from the target points, so a bipartite graph is given as (x_source, x_target) and (pos_source, pos_target) pairs.

Parameters:

channels (int) –

Feature width of the queries, keys and values.
num_groups (int) –

Number of channel groups sharing one attention weight.
attn_drop (float, default: 0.0 ) –

Dropout probability on the attention weights.
qkv_bias (bool, default: True ) –

Whether the query, key and value projections carry a bias.
pe_multiplier (bool, default: False ) –

Whether to scale the query-key relation by an encoding of the relative position.
pe_bias (bool, default: True ) –

Whether to shift the query-key relation and the values by an encoding of the relative position.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra activation arguments.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.
**kwargs (Unpack[MessagePassingParams], default: {} ) –

Extra arguments for MessagePassing.

KPConv [source] ¶

KPConv(
    spatial_dim: int,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    kp_radius: float,
    kp_sigma: float,
    kp_influence: str = "linear",
    fixed_position: Literal[
        "none", "center", "vertical"
    ] = "center",
    aggregation_mode: str = "sum",
    deformable: bool = False,
    modulated: bool = False,
    bias: bool = False,
    track_running_stats: bool = True,
)

Bases: MessagePassing

Kernel point convolution over a neighborhood graph.

Each of the \(K\) kernel points carries its own weight matrix, and a neighbor contributes to a kernel point with a weight given by their distance through the influence function. When deformable is set, a nested KPConv predicts a per-point offset for every kernel point.

edge_index lists [source, target] pairs: messages flow from the support (source) points to the query (target) points, and a bipartite graph is given as a (pos_source, pos_target) pair of positions.

The message is the \((E, K)\) influence of every kernel point on every edge; the aggregation pools the neighbor features per kernel point and only then applies that kernel point's weight matrix, so the \(K\) matrix products run over the \(N_t\) target points rather than over the \(E\) edges: the neighborhoods are laid out as a padded \((N_t, k)\) table and the pooling is one batched matrix product over it.

Parameters:

spatial_dim (int) –

Spatial dimension of the input point cloud.
in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
kernel_size (int) –

Number of kernel points \(K\).
kp_radius (float) –

Radius of the sphere the kernel points are placed on.
kp_sigma (float) –

Kernel extent, the distance over which a kernel point still influences a neighbor.
kp_influence (str, default: 'linear' ) –

Influence function: "constant", "linear", or "gaussian".
fixed_position (Literal['none', 'center', 'vertical'], default: 'center' ) –

Which kernel point is pinned: "none", "center", or "vertical".
aggregation_mode (str, default: 'sum' ) –

"sum" over all kernel points, or "closest" to keep only the nearest one.
deformable (bool, default: False ) –

Whether to predict per-point kernel offsets.
modulated (bool, default: False ) –

Whether the deformable branch also predicts a per-kernel-point modulation.
bias (bool, default: False ) –

Whether to add a bias to the output.
track_running_stats (bool, default: True ) –

Whether to keep the deformable statistics of the last forward pass (running_min_d2, running_deformed_kernel, running_offset_features) for losses.kpconv_deform_regularizer.

Shape

Input: \((N_s, C_\text{in})\) features, \((N, D)\) positions or a \(((N_s, D), (N_t, D))\) pair, \((2, E)\) edge index Output: \((N_t, C_\text{out})\) features

Methods:

configure_offsets –

Builds the rigid KPConv and the bias predicting the deformable offsets (and modulations).
configure_kernel –

Builds the fixed kernel point positions registered as the kernel buffer.

Attributes:

deformable (bool) –

Whether the kernel points are shifted by predicted per-point offsets.

deformable `property` ¶

deformable: bool

Whether the kernel points are shifted by predicted per-point offsets.

configure_offsets [source] ¶

configure_offsets() -> Tuple[Module, Parameter]

Builds the rigid KPConv and the bias predicting the deformable offsets (and modulations).

configure_kernel [source] ¶

configure_kernel() -> Tensor

Builds the fixed kernel point positions registered as the kernel buffer.

KPConvBlock [source] ¶

KPConvBlock(
    spatial_dim: int,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    kp_radius: float,
    kp_sigma: float,
    kp_influence: str = "linear",
    fixed_position: Literal[
        "none", "center", "vertical"
    ] = "center",
    aggregation_mode: str = "sum",
    deformable: bool = False,
    modulated: bool = False,
    act: Union[str, Callable, None] = "leaky_relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
)

Bases: Module

Kernel point convolution followed by normalization and activation.

KPResidualBlock [source] ¶

KPResidualBlock(
    spatial_dim: int,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    kp_radius: float,
    kp_sigma: float,
    kp_influence: str = "linear",
    fixed_position: Literal[
        "none", "center", "vertical"
    ] = "center",
    aggregation_mode: str = "sum",
    deformable: bool = False,
    modulated: bool = False,
    strided: bool = False,
    act: Union[str, Callable, None] = "leaky_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,
)

Bases: Module

Bottleneck residual block: a channel-reducing MLP, a KPConvBlock, and a channel-restoring MLP.

Set strided=True when the block maps a support cloud onto a coarser query cloud, so that the skip connection is max-pooled over the same neighborhoods.

LinearBlock [source] ¶

LinearBlock(
    in_channels: int,
    out_channels: int,
    act: Union[str, Callable, None] = None,
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    bias: bool = True,
    norm: Union[str, Callable, None] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Linear block consisting of a linear layer, normalization and activation.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
bias (bool, default: True ) –

Whether to use a bias for the linear layer.
act (Union[str, Callable, None], default: None ) –

Activation function to use. If None, no activation is applied.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra 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: None ) –

Normalization layer to use. If None, no normalization is applied.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the normalization layer.

Example

import torch

from torch_pointcloud.layers import LinearBlock

block = LinearBlock(64, 128, act="relu", norm="batch_norm", bias=False)
x = torch.randn(32, 64)
y = block(x)
print(y.shape)

OctreeAttention [source] ¶

OctreeAttention(
    channels: int,
    patch_size: int,
    num_heads: int,
    dilation: int = 1,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_rpe: bool = True,
)

Bases: Module

Multi-head self-attention restricted to windows of patch_size consecutive octree nodes.

A dilation above \(1\) interleaves the nodes before windowing, so a patch spans a wider region at the same cost. The attention masks and relative positions come from the OctreeT, which must have been built with construct_all_attention_context().

Parameters:

channels (int) –

Number of input and output channels.
patch_size (int) –

Number of octree nodes attending to each other.
num_heads (int) –

Number of attention heads.
dilation (int, default: 1 ) –

Stride between the nodes of a patch.
qkv_bias (bool, default: True ) –

Whether to use a bias in the QKV projection.
qk_scale (Optional[float], default: None ) –

Scaling factor for the QK matrix. Defaults to the inverse square root of the head dimension.
attn_drop (float, default: 0.0 ) –

Dropout rate for the attention.
proj_drop (float, default: 0.0 ) –

Dropout rate for the output projection.
use_rpe (bool, default: True ) –

Whether to add the relative position encoding to the attention logits.

Methods:

forward_attn –

Runs the windowed attention on patched features of shape \((N, K, C)\), returning \((N \cdot K, C)\).
forward_proj –

Applies the output projection and its dropout.
forward_rpe –

Adds the relative position bias to the attention logits, or returns them unchanged when use_rpe is off.

Attributes:

dilated (bool) –

Whether the patches interleave nodes rather than taking them consecutively.

dilated `property` ¶

dilated: bool

Whether the patches interleave nodes rather than taking them consecutively.

forward_attn [source] ¶

forward_attn(
    x: Tensor, rel_pos: Tensor, mask: Tensor
) -> Tensor

Runs the windowed attention on patched features of shape \((N, K, C)\), returning \((N \cdot K, C)\).

forward_proj [source] ¶

forward_proj(x: Tensor) -> Tensor

Applies the output projection and its dropout.

forward_rpe [source] ¶

forward_rpe(
    attn: Tensor, rel_pos: Optional[Tensor] = None
) -> Tensor

Adds the relative position bias to the attention logits, or returns them unchanged when use_rpe is off.

OctreeRelativePositionEncoding [source] ¶

OctreeRelativePositionEncoding(
    patch_size: int, num_heads: int, dilation: int = 1
)

Bases: Module

Relative Position Encoding (RPE) module used within the OctreeAttention module.

Parameters:

patch_size (int) –

The patch size to use for the RPE.
num_heads (int) –

The number of heads to use for the RPE.
dilation (int, default: 1 ) –

The dilation to use for the RPE.

Methods:

pos_to_idx –

Clamps the relative positions to \(\pm\) pos_bnd and maps each axis to its slice of the RPE table.

pos_to_idx [source] ¶

pos_to_idx(pos: Tensor) -> Tensor

Clamps the relative positions to \(\pm\) pos_bnd and maps each axis to its slice of the RPE table.

OctreeT [source] ¶

OctreeT(
    depth: int,
    patch_size: int,
    dilation: int,
    full_depth: int = 2,
    batch_size: int = 1,
    device: Union[device, str] = "cpu",
    **kwargs: Any,
)

Bases: Octree

An enhanced Octree with transformer-specific capabilities (patching, dilation, masking).

Can be instantiated directly like a standard Octree, or created from an existing Octree instance using OctreeT.from_octree().

Once a OctreeT is instantiated, you can build the transformer context (i.e. all attention masks and relative positions) by calling the method construct_all_attention_context().

Example

>>> octree_t = OctreeT.from_octree(octree, patch_size=26, dilation=4)  # doctest: +SKIP
>>> octree_t.construct_all_attention_context(  # doctest: +SKIP
...     nempty=True,
...     min_depth=6,
...     max_depth=10,
... )
>>> octree_t.masks[6].shape  # doctest: +SKIP
>>> octree_t.dilated_masks[6].shape  # doctest: +SKIP
>>> octree_t.rel_pos[6].shape  # doctest: +SKIP
>>> octree_t.dilated_rel_pos[6].shape  # doctest: +SKIP

Methods:

reset –

Resets the OctreeT to its initial state.
from_octree –

Creates an OctreeT from an Octree.
construct_all_attention_context –

Constructs all attention context for the octree.
construct_attention_context –

Calculates attention masks, relative positions, and padding indices
pad_to_patch_size –

Pads x along its first dimension so the node count at depth is a whole number of patches.
unpad –

Drops the padding added by pad_to_patch_size, restoring the real node count at depth.

Attributes:

block_size (int) –

Number of octree nodes a dilated patch spans, \(\text{patch\_size} \cdot \text{dilation}\).

block_size `property` ¶

block_size: int

Number of octree nodes a dilated patch spans, \(\text{patch\_size} \cdot \text{dilation}\).

reset [source] ¶

reset() -> None

Resets the OctreeT to its initial state.

from_octree `classmethod` [source] ¶

from_octree(
    octree: Octree,
    patch_size: int,
    dilation: int,
    **kwargs: Any,
) -> OctreeT

Creates an OctreeT from an Octree.

Parameters:

octree (Octree) –

The Octree to create the OctreeT from.
patch_size (int) –

The patch size to use for the OctreeT.
dilation (int) –

The dilation to use for the OctreeT.
**kwargs (Any, default: {} ) –

Additional keyword arguments to pass to the OctreeT constructor.

construct_all_attention_context [source] ¶

construct_all_attention_context(
    nempty: bool = False,
    min_depth: Optional[int] = None,
    max_depth: Optional[int] = None,
) -> None

Constructs all attention context for the octree.

Parameters:

nempty (bool, default: False ) –

Whether to use non-empty nodes.
min_depth (Optional[int], default: None ) –

The start depth of the octree to construct the context for.
max_depth (Optional[int], default: None ) –

The end depth of the octree to construct the context for.

construct_attention_context [source] ¶

construct_attention_context(
    depth: int, nempty: bool = False
) -> None

Calculates attention masks, relative positions, and padding indices required for attention operations.

Parameters:

depth (int) –

The depth of the octree to construct the context for.
nempty (bool, default: False ) –

Whether to use non-empty nodes.

pad_to_patch_size [source] ¶

pad_to_patch_size(
    x: Tensor, depth: int, fill_value: float = 0
) -> Tensor

Pads x along its first dimension so the node count at depth is a whole number of patches.

unpad [source] ¶

unpad(x: Tensor, depth: int) -> Tensor

Drops the padding added by pad_to_patch_size, restoring the real node count at depth.

OctreeConvBlock [source] ¶

OctreeConvBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: Union[int, Sequence[int]],
    stride: int = 1,
    nempty: bool = False,
    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,
    method: str = "explicit_gemm",
    max_buffer: int = MAX_BUFFER,
)

Bases: Module

Octree convolution followed by normalization and activation.

Wraps ocnn.nn.OctreeConv with a norm / act pair built by create_norm / create_act. With act_first=True the activation runs before the normalization instead of after.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
kernel_size (Union[int, Sequence[int]]) –

Convolution kernel size (an int is broadcast to all axes).
stride (int, default: 1 ) –

Convolution stride; 2 downsamples the octree by one depth level.
nempty (bool, default: False ) –

Whether the features only cover non-empty octree nodes.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.
act_first (bool, default: False ) –

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

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

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

Whether the convolution uses a bias.
method (str, default: 'explicit_gemm' ) –

ocnn convolution implementation (e.g. "explicit_gemm").
max_buffer (int, default: MAX_BUFFER ) –

Maximum buffer size (in elements) used by the ocnn convolution.

Shape

Input: \((M_\text{in}, C_\text{in})\) octree features at depth.
Output: \((M_\text{out}, C_\text{out})\) octree features (at depth - 1 when stride=2).

OctreeDeconvBlock [source] ¶

OctreeDeconvBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: Union[int, Sequence[int]],
    stride: int = 1,
    nempty: bool = False,
    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,
    method: str = "explicit_gemm",
    max_buffer: int = MAX_BUFFER,
)

Bases: Module

Octree transposed convolution followed by normalization and activation.

Wraps ocnn.nn.OctreeDeconv with a norm / act pair built by create_norm / create_act. With act_first=True the activation runs before the normalization instead of after.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
kernel_size (Union[int, Sequence[int]]) –

Convolution kernel size (an int is broadcast to all axes).
stride (int, default: 1 ) –

Convolution stride; 2 upsamples the octree by one depth level.
nempty (bool, default: False ) –

Whether the features only cover non-empty octree nodes.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.
act_first (bool, default: False ) –

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

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

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

Whether the convolution uses a bias.
method (str, default: 'explicit_gemm' ) –

ocnn convolution implementation (e.g. "explicit_gemm").
max_buffer (int, default: MAX_BUFFER ) –

Maximum buffer size (in elements) used by the ocnn convolution.

Shape

Input: \((M_\text{in}, C_\text{in})\) octree features at depth.
Output: \((M_\text{out}, C_\text{out})\) octree features (at depth + 1 when stride=2).

PDNorm [source] ¶

PDNorm(
    channels: int,
    conditions: Sequence[str],
    norm: Union[str, Callable, None] = "batch_norm",
    decouple: bool = True,
    *,
    dim: int = 1,
    **norm_kwargs: Any,
)

Bases: Module

Prompt-driven normalization that routes each batch through a per-condition norm.

Multi-dataset joint training feeds batches drawn from a single dataset at a time, identified by a string condition. When decouple is True, PDNorm holds one independent inner norm per condition (an nn.ModuleList indexed by conditions.index(condition)), so each dataset keeps its own running statistics and affine parameters. When decouple is False, a single shared norm is applied regardless of the condition.

Inner norms are built with create_norm, so any name / class / instance accepted there is valid. The decoupled layout stores children under norm.{i} keys, matching the order of conditions.

Introduced in Towards Large-scale 3D Representation Learning with Multi-dataset Point Prompt Training

Parameters:

channels (int) –

Number of feature channels.
conditions (Sequence[str]) –

Ordered condition names; index \(i\) selects norm.{i} when decoupled.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Inner norm passed to create_norm (name, class, instance, or None).
decouple (bool, default: True ) –

If True, use one norm per condition; if False, share a single norm.
dim (int, default: 1 ) –

Spatial dimensionality hint forwarded to each inner create_norm (see its dim argument).
**norm_kwargs (Any, default: {} ) –

Extra keyword arguments forwarded to each inner norm constructor.

Shape

Input: \((N, C)\) packed features with \(C =\) channels.
Output: \((N, C)\), same shape as the input.

Example

import torch
from torch_pointcloud.layers import PDNorm

norm = PDNorm(64, conditions=["ScanNet", "S3DIS"], norm="batch_norm")
x = torch.randn(32, 64)
y = norm(x, condition="S3DIS")
print(y.shape)

PointPatchEmbed [source] ¶

PointPatchEmbed(
    embed_dim: int,
    in_channels: int = 3,
    local_channels: Sequence[int] = (128, 256),
    global_channels: Sequence[int] = (512,),
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Mini-PointNet patch (token) embedding for grouped point clouds.

Embeds each local group of points into a single token via a two-stage shared MLP with an intermediate per-group max-pool and global-feature concatenation, as used by the masked / autoregressive point self-supervised models (Point-MAE, Point-BERT, PointGPT, Point-M2AE). A shared \(1 \times 1\) convolution over \((B, C, M)\) is equivalent to a MLP over the feature dim, so both stages are plain PyG MLPs. local_mlp maps in_channels \(\to \text{local\_channels}\); the per-group max-pool is concatenated to give \(2 \cdot \text{local\_channels}[-1]\); global_mlp maps that to embed_dim through global_channels. The module is permutation-invariant over the points within a group.

Parameters:

embed_dim (int) –

Output token dimension.
in_channels (int, default: 3 ) –

Channels per input point (\(3\) for coordinates only, plus any concatenated features).
local_channels (Sequence[int], default: (128, 256) ) –

Hidden widths of the per-point MLP (in_channels is prepended).
global_channels (Sequence[int], default: (512,) ) –

Hidden widths of the per-group MLP (\(2 \cdot \text{local\_channels}[-1]\) input and embed_dim output are added).
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra arguments for the activation.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization of the MLPs.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the normalization.

Shape

Input: \((B, G, M, C)\) where \(B\) is the batch size, \(G\) the number of groups, \(M\) the group size, and \(C\) = in_channels.
Output: \((B, G, D)\) where \(D\) = embed_dim.

Example

import torch
from torch_pointcloud.layers import PointPatchEmbed

embed = PointPatchEmbed(embed_dim=384)
tokens = embed(torch.randn(2, 64, 32, 3))
print(tokens.shape)

PointTransformerConv [source] ¶

PointTransformerConv(
    in_channels: Union[int, Tuple[int, int]],
    out_channels: int,
    spatial_dim: int = 3,
    num_groups: int = 8,
    pos_nn: Optional[Callable[[Tensor], Tensor]] = None,
    attn_nn: Optional[Callable[[Tensor], Tensor]] = None,
    add_self_loops: bool = False,
    bias: bool = False,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

The Point Transformer layer from the "Point Transformer" paper by Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, Vladlen Koltun.

Note

This implementation was adapted from the PyTorch Geometric library, and supports the num_groups parameter to behave like the original implementation.

\[ \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \alpha_{i,j} \left(\mathbf{W}_3 \mathbf{x}_j + \delta_{ij} \right), \]

where the attention coefficients \(\alpha_{i,j}\) and positional embedding \(\delta_{ij}\) are computed as

\[ \alpha_{i,j}= \textrm{softmax} \left( \gamma_\mathbf{\Theta} (\mathbf{W}_1 \mathbf{x}_i - \mathbf{W}_2 \mathbf{x}_j + \delta_{i,j}) \right) \]

and

\[ \delta_{i,j}= h_{\mathbf{\Theta}}(\mathbf{p}_i - \mathbf{p}_j), \]

with \(\gamma_\mathbf{\Theta}\) and \(h_\mathbf{\Theta}\) denoting neural networks, i.e. MLPs, and \(\mathbf{P} \in \mathbb{R}^{N \times D}\) defines the position of each point.

Parameters:

in_channels (int or tuple) –

Size of each input sample, or -1 to derive the size from the first input(s) to the forward method. A tuple corresponds to the sizes of source and target dimensionalities.
out_channels (int) –

Size of each output sample.
pos_nn (Module, default: None ) –

A neural network \(h_\mathbf{\Theta}\) which maps relative spatial coordinates pos_j - pos_i of shape \([-1, 3]\) to shape \([-1, \text{out\_channels}]\). Will default to a torch.nn.Linear transformation if not further specified.
attn_nn (Module, default: None ) –

A neural network \(\gamma_\mathbf{\Theta}\) which maps transformed node features of shape \([-1, \text{out\_channels}]\) to shape \([-1, \text{out\_channels}]\).
add_self_loops (bool, default: False ) –

If False, do not add self-loops to the input graph.
bias (bool, default: False ) –

Whether the query, key and value projections carry a bias.

Shapes

input: node features \((|\mathcal{V}|, F_{in})\) or \(((|\mathcal{V_s}|, F_{s}), (|\mathcal{V_t}|, F_{t}))\) if bipartite, positions \((|\mathcal{V}|, 3)\) or \(((|\mathcal{V_s}|, 3), (|\mathcal{V_t}|, 3))\) if bipartite, edge indices \((2, |\mathcal{E}|)\)
output: node features \((|\mathcal{V}|, F_{out})\) or \(((|\mathcal{V}_t|, F_{out}))\) if bipartite

PointConv [source] ¶

PointConv(
    local_nn: Module,
    weight_nn: Module,
    add_self_loops: bool = False,
    eps: float = 1e-06,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

PointConv message passing: a per-edge feature outer-producted with a learned continuous weight.

local_nn embeds the relative position concatenated with the neighbor feature, weight_nn turns the relative position alone into the convolution weights, and their outer product is summed over the neighbors of each point.

Parameters:

local_nn (Module) –

Network embedding the relative positions and neighbor features.
weight_nn (Module) –

Network mapping a relative position to the convolution weights.
add_self_loops (bool, default: False ) –

Whether to add a self loop to every point before propagating.
eps (float, default: 1e-06 ) –

Numerical stability constant.
**kwargs (Unpack[MessagePassingParams], default: {} ) –

Extra arguments for MessagePassing.

PointConvDensity [source] ¶

PointConvDensity(
    local_nn: Module,
    weight_nn: Module,
    density_nn: Module,
    add_self_loops: bool = False,
    eps: float = 1e-06,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

PointConv with inverse density reweighting.

Each neighbor's embedded feature is scaled by density_nn applied to its density relative to the densest neighbor, so that densely sampled regions do not dominate the sum.

Parameters:

local_nn (Module) –

Network embedding the relative positions and neighbor features.
weight_nn (Module) –

Network mapping a relative position to the convolution weights.
density_nn (Module) –

Network mapping a relative density to a per-edge scale.
add_self_loops (bool, default: False ) –

Whether to add a self loop to every point before propagating.
eps (float, default: 1e-06 ) –

Numerical stability constant.
**kwargs (Unpack[MessagePassingParams], default: {} ) –

Extra arguments for MessagePassing.

PointConvDensityGlobalSetAbstraction [source] ¶

PointConvDensityGlobalSetAbstraction(
    in_channels: int,
    channels: Sequence[int],
    bandwidth: float,
    weight_channels: Sequence[int] = (8, 8),
    density_channels: Sequence[int] = (16, 8),
    expansion: int = 16,
    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,
    pool: PoolLike = "mean",
    spatial_dim: int = 3,
)

Bases: Module

Global PointConv set-abstraction block with inverse-density re-weighting.

Parameters:

in_channels (int) –

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

Per-layer channel sizes of the feature MLP.
bandwidth (float) –

Bandwidth of the Gaussian kernel density estimate.
weight_channels (Sequence[int], default: (8, 8) ) –

Hidden channel sizes of the weight net applied to relative positions.
density_channels (Sequence[int], default: (16, 8) ) –

Hidden channel sizes of the density net.
expansion (int, default: 16 ) –

Channel expansion factor of the final matrix multiplication.
pool (PoolLike, default: 'mean' ) –

Pooling used to place the single output position of each sample.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.

PointConvDensitySetAbstraction [source] ¶

PointConvDensitySetAbstraction(
    in_channels: int,
    num_neighbors: int,
    channels: Sequence[int],
    bandwidth: float = 1.0,
    weight_channels: Sequence[int] = (8, 8),
    density_channels: Sequence[int] = (16, 8),
    expansion: int = 16,
    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,
    spatial_dim: int = 3,
    downsample: Optional[Module] = None,
)

Bases: Module

PointConv set-abstraction block with inverse-density re-weighting.

Same layout as PointConvSetAbstraction, with a Gaussian kernel density estimated per point; the inverse density is transformed by a density net and re-weights the grouped features.

Parameters:

in_channels (int) –

Number of input feature channels.
num_neighbors (int) –

Number of neighbors gathered per output point.
channels (Sequence[int]) –

Per-layer channel sizes of the feature MLP.
bandwidth (float, default: 1.0 ) –

Bandwidth of the Gaussian kernel density estimate.
weight_channels (Sequence[int], default: (8, 8) ) –

Hidden channel sizes of the weight net applied to relative positions.
density_channels (Sequence[int], default: (16, 8) ) –

Hidden channel sizes of the density net.
expansion (int, default: 16 ) –

Channel expansion factor of the final matrix multiplication.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.
downsample (Optional[Module], default: None ) –

Optional module returning the sampled indices (e.g. FPS). If None, the resolution is unchanged.

PointConvGlobalSetAbstraction [source] ¶

PointConvGlobalSetAbstraction(
    in_channels: int,
    channels: Sequence[int],
    weight_channels: Sequence[int] = (8, 8),
    expansion: int = 16,
    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,
    aggr: PoolLike = "mean",
    spatial_dim: int = 3,
)

Bases: Module

Global PointConv set-abstraction block: one weight-net convolution over each whole sample.

Parameters:

in_channels (int) –

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

Per-layer channel sizes of the feature MLP.
weight_channels (Sequence[int], default: (8, 8) ) –

Hidden channel sizes of the weight net applied to relative positions.
expansion (int, default: 16 ) –

Channel expansion factor of the final matrix multiplication.
aggr (PoolLike, default: 'mean' ) –

Pooling used to place the single output position of each sample.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.

PointConvSetAbstraction [source] ¶

PointConvSetAbstraction(
    in_channels: int,
    num_neighbors: int,
    channels: Sequence[int],
    weight_channels: Sequence[int] = (8, 8),
    expansion: int = 16,
    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,
    spatial_dim: int = 3,
    downsample: Optional[Module] = None,
)

Bases: Module

PointConv set-abstraction block: optional downsampling, \(k\)-NN grouping, and a weight-net continuous convolution.

Parameters:

in_channels (int) –

Number of input feature channels.
num_neighbors (int) –

Number of neighbors gathered per output point.
channels (Sequence[int]) –

Per-layer channel sizes of the feature MLP.
weight_channels (Sequence[int], default: (8, 8) ) –

Hidden channel sizes of the weight net applied to relative positions.
expansion (int, default: 16 ) –

Channel expansion factor of the final matrix multiplication.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.
downsample (Optional[Module], default: None ) –

Optional module returning the sampled indices (e.g. FPS). If None, the resolution is unchanged.

PointNet2Conv [source] ¶

PointNet2Conv(
    local_nn: Module,
    add_self_loops: bool = True,
    use_pos: bool = True,
    pos_first: bool = False,
    pos_scale: Optional[float] = None,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

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

Each message combines the neighbor features \(x_j\) with the relative position \(p_j - p_i\) and applies local_nn; the messages of each centroid are then aggregated (aggr, max by default).

Parameters:

local_nn (Module) –

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

Whether to add self-loops to the edge index.
use_pos (bool, default: True ) –

Concatenate the relative position \(p_j - p_i\) to the neighbor features.
pos_first (bool, default: False ) –

Concatenate the relative position before the features (cat([p_j - p_i, x_j])) instead of after.
pos_scale (Optional[float], default: None ) –

Divide the relative position by this value, e.g. the ball-query radius to normalize it.
**kwargs (Unpack[MessagePassingParams], default: {} ) –

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

PointNet2FeaturePropagation [source] ¶

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",
    eps: float = 1e-16,
)

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

Parameters:

channels (Sequence[int]) –

Per-layer channel sizes of the MLP, starting with its input width (interpolated plus skip channels).
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).
dropout (float, default: 0.0 ) –

Dropout rate inside the MLP.
bias (Union[bool, List[bool]], default: True ) –

Whether the MLP linear layers use a bias.
plain_last (bool, default: True ) –

Leave the last MLP layer without normalization and activation.
weighting (Literal['squared', 'inverse'], default: 'inverse' ) –

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.

PointNet2GlobalSetAbstraction [source] ¶

PointNet2GlobalSetAbstraction(
    in_channels: int,
    channels: Sequence[int],
    spatial_dim: int = 3,
    dropout: float = 0.0,
    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,
    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.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.
dropout (float, default: 0.0 ) –

Dropout rate inside the MLP.
bias (bool, default: False ) –

Whether the MLP linear layers use a bias.
use_pos (bool, default: False ) –

Concatenate the absolute point positions to x before the MLP. Unlike PointNet2SetAbstraction 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.
aggr (PoolLike, default: 'max' ) –

Pooling operation applied per batch element.

PointNet2SetAbstraction [source] ¶

PointNet2SetAbstraction(
    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,
    dropout: float = 0.0,
    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,
    aggr: AggrType = "max",
    sort_neighbors: bool = False,
)

Bases: Module

Set-abstraction block (PointNet++ SSG / MSG) 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:

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 (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).
radii (Union[float, Sequence[float]]) –

Ball-query radius per scale.
num_neighbors (Union[int, Sequence[int]]) –

Maximum number of neighbors per scale.
spatial_dim (int, default: 3 ) –

Dimension of point coordinates.
dropout (float, default: 0.0 ) –

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

Whether the MLP linear layers use a bias.
use_pos (bool, default: True ) –

Concatenate the relative position to the grouped features.
normalize_pos (bool, default: True ) –

Divide the relative position by the ball-query radius.
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.
aggr (AggrType, default: 'max' ) –

Aggregation of the messages of each centroid.
sort_neighbors (bool, default: False ) –

Keep the smallest source indices when a ball holds more than num_neighbors points, as the reference query_ball_point does.

PointNeXtConv [source] ¶

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

PointNeXtResidualBlock [source] ¶

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.

PointNeXtSetAbstraction [source] ¶

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

CatPool [source] ¶

CatPool(
    pools: Sequence[PoolLike] = ("max", "mean"),
    dim: int = 0,
    dim_size: Optional[int] = None,
)

Bases: Module

Runs several pools on the same input and concatenates their outputs along the feature dim.

Parameters:

pools (Sequence[PoolLike], default: ('max', 'mean') ) –

Pools to combine, each resolved by create_pool (a name, class, or instance).
dim (int, default: 0 ) –

Dimension along which each pool reduces.
dim_size (Optional[int], default: None ) –

Number of output segments \(B\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C \cdot P)\) where \(P\) is the number of pools.

Example

import torch
from torch_pointcloud.layers import CatPool

pool = CatPool(pools=("max", "mean"))
x = torch.randn(6, 4)
batch = torch.tensor([0, 0, 0, 1, 1, 1])
out = pool(x, batch)  # (2, 8)

Attributes:

num_pools (int) –

Number of pools \(P\) concatenated, i.e. the feature multiplier.

num_pools `property` ¶

num_pools: int

Number of pools \(P\) concatenated, i.e. the feature multiplier.

MaxPool [source] ¶

MaxPool(dim: int = 0, dim_size: Optional[int] = None)

Bases: Module

Per-segment max pooling over a packed batch, via torch_geometric.utils.scatter(reduce="max").

Parameters:

dim (int, default: 0 ) –

Dimension along which to pool.
dim_size (Optional[int], default: None ) –

Number of output segments \(B\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

MeanPool [source] ¶

MeanPool(dim: int = 0, dim_size: Optional[int] = None)

Bases: Module

Per-segment mean pooling over a packed batch, via torch_geometric.utils.scatter(reduce="mean").

Parameters:

dim (int, default: 0 ) –

Dimension along which to pool.
dim_size (Optional[int], default: None ) –

Number of output segments \(B\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

MinPool [source] ¶

MinPool(dim: int = 0, dim_size: Optional[int] = None)

Bases: Module

Per-segment min pooling over a packed batch, via torch_geometric.utils.scatter(reduce="min").

Parameters:

dim (int, default: 0 ) –

Dimension along which to pool.
dim_size (Optional[int], default: None ) –

Number of output segments \(B\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

MulPool [source] ¶

MulPool(dim: int = 0, dim_size: Optional[int] = None)

Bases: Module

Per-segment product pooling over a packed batch, via torch_geometric.utils.scatter(reduce="mul").

Parameters:

dim (int, default: 0 ) –

Dimension along which to pool.
dim_size (Optional[int], default: None ) –

Number of output segments \(B\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

SumPool [source] ¶

SumPool(dim: int = 0, dim_size: Optional[int] = None)

Bases: Module

Per-segment sum pooling over a packed batch, via torch_geometric.utils.scatter(reduce="sum").

Parameters:

dim (int, default: 0 ) –

Dimension along which to pool.
dim_size (Optional[int], default: None ) –

Number of output segments \(B\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

PVConv [source] ¶

PVConv(
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    resolution: int,
    use_se: bool = False,
    normalize: bool = True,
    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,
)

Bases: Module

Point-voxel convolution: a dense 3D conv branch summed with a per-point MLP branch.

The voxel branch captures the neighborhood context at a coarse resolution and is devoxelized back with trilinear interpolation, while the point branch keeps the fine per-point detail.

Shape

Input: \((N, C_\text{in})\) features, \((N, 3)\) positions, \((N,)\) batch index Output: \((N, C_\text{out})\)

Parameters:

in_channels (int) –

Input channel count.
out_channels (int) –

Output channel count.
kernel_size (int) –

Kernel size of the voxel convolutions.
resolution (int) –

Voxel grid resolution \(R\) along each axis.
use_se (bool, default: False ) –

Whether to gate the voxel branch with an SE3d block.
normalize (bool, default: True ) –

Whether to rescale each cloud to the unit grid before voxelizing.
act (Union[str, Callable, None], default: 'relu' ) –

Activation, name resolved by create_act. None disables.
act_first (bool, default: False ) –

If True, run activation before normalization.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the activation.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization, name resolved by create_norm. None disables.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the normalization.

SE3d [source] ¶

SE3d(
    channels: int,
    reduction: int = 8,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Squeeze-and-excitation gate for dense voxel grids.

Global-average-pools each channel, passes it through a bottleneck MLP, and rescales the grid by the resulting per-channel gate.

Shape

Input: \((B, C, R, R, R)\) Output: \((B, C, R, R, R)\)

Parameters:

channels (int) –

Number of channels \(C\).
reduction (int, default: 8 ) –

Bottleneck ratio of the squeeze layer.
act (Union[str, Callable, None], default: 'relu' ) –

Activation between the squeeze and excitation layers, name resolved by create_act.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the activation.

Voxelization [source] ¶

Voxelization(resolution: int, normalize: bool = True)

Bases: Module

Averages packed point features into a dense voxel grid.

Positions are centered per cloud and, when normalize is set, rescaled by the cloud's largest radius so every cloud fills the grid. The continuous grid coordinates are returned alongside the voxels, for trilinear devoxelization.

Shape

Input: \((N, C)\) features, \((N, 3)\) positions, \((N,)\) batch index Output: \((B, C, R, R, R)\) voxels, \((N, 3)\) grid coordinates

Parameters:

resolution (int) –

Grid resolution \(R\) along each axis.
normalize (bool, default: True ) –

Whether to rescale each cloud to the unit grid before voxelizing.

AttentivePooling [source] ¶

AttentivePooling(
    in_channels: int,
    out_channels: int,
    *,
    act: Union[str, Callable, None],
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None],
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

Attention-weighted aggregation of neighbor features.

Each edge concatenates its source feature with a per-edge spatial encoding, scores the result with a no-bias linear layer, softmax-normalizes the scores over the neighbors of each target point, sums the score-weighted features per target, then projects them with a Linear+norm+act block.

edge_index lists [source, target] pairs and edge_attr holds the \((E, C_\text{pos})\) spatial encoding of each edge.

Parameters:

in_channels (int) –

Channels of each edge feature, the source feature and the spatial encoding concatenated.
out_channels (int) –

Output channels after the post-aggregation MLP.
**kwargs (Unpack[MessagePassingParams], default: {} ) –

Extra arguments for MessagePassing.

LocalFeatureAggregation [source] ¶

LocalFeatureAggregation(
    d_out: int,
    *,
    act: Union[str, Callable, None],
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None],
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
)

Bases: Module

Local feature aggregation: two rounds of spatial encoding and attentive pooling.

Each round grows the receptive field, taking a per-point feature of width \(d_\text{out} / 2\) to \(d_\text{out}\). The 10-channel relative positional encoding keeps the channel order of QingyongHu/RandLA-Net (cat([rel_dist, rel_xyz, xyz_i, xyz_j], dim=-1)) so pretrained weights load without permuting the first kernel, and the second encoding re-projects the output of the first.

LocalSpatialEncoding [source] ¶

LocalSpatialEncoding(
    in_channels: int,
    out_channels: int,
    *,
    act: Union[str, Callable, None],
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None],
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
)

Bases: Module

Per-edge spatial encoding MLP.

Wraps a single Linear+norm+act block that lifts an input feature to out_channels. Used twice per LocalFeatureAggregation: first on the raw 10-channel relative positional encoding, then on its output.

RandLANetResidualBlock [source] ¶

RandLANetResidualBlock(
    d_in: int,
    d_out: int,
    num_neighbors: int,
    *,
    act: Union[str, Callable, None],
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None],
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
)

Bases: Module

RandLA-Net dilated residual block.

Maps d_in channels to \(2 \cdot d_\text{out}\) via a residual path of MLP -> LocalFeatureAggregation -> MLP plus a parallel Linear+norm shortcut. mlp2 and shortcut carry no activation; the configured activation is applied once after the residual sum.

Parameters:

d_in (int) –

Number of input channels.
d_out (int) –

"Configuration" channel count; the block actually outputs \(2 \cdot d_\text{out}\).
num_neighbors (int) –

Number of neighbors for the local feature aggregation.

Point3DRoPE [source] ¶

Point3DRoPE(head_dim: int, base: float = 10.0)

Bases: Module

3D Rotary Position Embedding for point cloud attention.

Parameters:

head_dim (int) –

Channel dimension of a single attention head. Must be divisible by 6: the head splits into three axis chunks, each of which needs an even size for the half-rotation.
base (float, default: 10.0 ) –

RoPE frequency base (\(\theta\)). Smaller values encode finer spatial detail.

Inputs

q: Query tensor of shape \((N, H, D)\) where \(D\) is head_dim. k: Key tensor of the same shape as q. pos: Real-valued 3D positions of shape \((N, 3)\) corresponding to each token.

Outputs

Tuple (q_rot, k_rot) of tensors with the same shapes as q and k.

PatchLayout [source] ¶

Bases: TypedDict

Padding and patch offsets of one point set for serialized attention.

Every block of a stage attends over the same points, so the encoder builds this once per stage with patch_layout and hands it to its blocks instead of recomputing it in every attention call.

Attributes:

padded_indices (Tensor) –

Row of every padded slot, shape \((N_p,)\).
unpadded_indices (Tensor) –

Padded slot of every row, shape \((N,)\).
padded_batch (Tensor) –

Batch index of every padded slot, shape \((N_p,)\).
cu_seqlens (Tensor) –

Cumulative patch lengths as int32, shape \((P + 1,)\), the layout flash attention takes.

RelativePositionalEncoding [source] ¶

RelativePositionalEncoding(patch_size: int, num_heads: int)

Bases: Module

Table-based relative position bias added to the attention logits of a serialized patch.

Each axis of the clamped relative grid coordinates indexes its own slice of a shared table, and the three per-head biases are summed. Clamping leaves \(\text{rpe\_num}^3\) distinct offsets, so the three rows are summed once per combination into a lookup table and the bias is a single gather. Summing them per point pair instead would hold \((\text{pairs}, 3)\) offsets and \((\text{pairs}, 3, H)\) gathered rows, which on a large cloud is the peak allocation of the whole model.

Parameters:

patch_size (int) –

Number of points attending to each other, setting the clamping boundary.
num_heads (int) –

Number of attention heads.

Shape

Input: \((P, K, 3)\) integer grid coordinates, \(K\) points per patch.
Output: \((P, H, K, K)\) bias, aligned with the attention logits.

Example

import torch

from torch_pointcloud.layers import RelativePositionalEncoding

rpe = RelativePositionalEncoding(patch_size=8, num_heads=4)
bias = rpe(torch.randint(0, 16, (2, 8, 3)))
bias.shape  # (2, 4, 8, 8)

SerializedAttention [source] ¶

SerializedAttention(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
)

Bases: Module

Vanilla serialized attention from Point Transformer V3.

No positional information is added inside attention itself: relative structure comes from the conditional position embedding (CPE) applied around each block.

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

Methods:

patch_layout –

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

patch_layout [source] ¶

patch_layout(batch: Tensor) -> PatchLayout

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

SerializedAttentionRoPE [source] ¶

SerializedAttentionRoPE(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
    rope_base: float = 10.0,
)

Bases: Module

Serialized attention with 3D rotary position embedding from Utonia.

Rotates \(Q\), \(K\) via Point3DRoPE using the real-valued metric position of each token. Flash Attention is supported and uses bfloat16 (matching upstream Utonia's reference implementation).

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

Methods:

patch_layout –

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

patch_layout [source] ¶

patch_layout(batch: Tensor) -> PatchLayout

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

SerializedAttentionRPE [source] ¶

SerializedAttentionRPE(
    channels: int,
    num_heads: int,
    patch_size: int,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
)

Bases: Module

Serialized attention with the relative position bias from PT-V3.

Adds a learned per-head bias indexed by the integer voxel-grid offset between query and key inside each patch. Flash Attention does not support arbitrary attention biases, so this variant always uses the manual softmax path.

Note

The input must be batch-blocked: each sample's points contiguous (a sorted batch), and any serialized_order must keep samples contiguous (the serialization codes embed the batch index). Otherwise patches silently mix points from different samples and attention leaks across the batch.

Methods:

patch_layout –

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

patch_layout [source] ¶

patch_layout(batch: Tensor) -> PatchLayout

The PatchLayout this layer needs for batch; build it once per stage and pass it to forward.

SerializedPool [source] ¶

SerializedPool(
    in_channels: int,
    out_channels: int,
    stride: int = 2,
    bias: bool = True,
    act: Union[str, Callable, None] = None,
    norm: Union[str, Callable, None] = None,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    reduce: Literal["sum", "mean", "min", "max"] = "max",
)

Bases: Module

Grid pooling driven by the serialization code: points sharing a coarser code are reduced to one.

Truncating the serialization code by \(3 \cdot \log_2(\text{stride})\) bits is exactly a grid subsampling, so no neighbor search is needed. Features are projected then reduced, positions are averaged, and the coarsened code is returned so the next stage can pool again.

Parameters:

in_channels (int) –

Input channel count.
out_channels (int) –

Output channel count.
stride (int, default: 2 ) –

Downsampling factor along each axis. Must be a power of \(2\).
bias (bool, default: True ) –

Whether the projection has a bias term.
act (Union[str, Callable, None], default: None ) –

Activation, name resolved by create_act. None disables.
norm (Union[str, Callable, None], default: None ) –

Normalization, name resolved by create_norm. None disables.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the activation.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the normalization.
reduce (Literal['sum', 'mean', 'min', 'max'], default: 'max' ) –

Reduction applied within a cell: "sum", "mean", "min", or "max".

SerializedUpsample [source] ¶

SerializedUpsample(
    in_channels: int,
    skip_channels: int,
    out_channels: int,
    norm: Union[str, Callable, None] = None,
    act: Union[str, Callable, None] = None,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
)

Bases: Module

Undoes a SerializedPool step: scatters the pooled features back and adds the projected skip.

Both branches get their own projection, normalization and activation before being summed.

Parameters:

in_channels (int) –

Channel count of the pooled features.
skip_channels (int) –

Channel count of the skip connection.
out_channels (int) –

Output channel count.
norm (Union[str, Callable, None], default: None ) –

Normalization, name resolved by create_norm. None disables.
act (Union[str, Callable, None], default: None ) –

Activation, name resolved by create_act. None disables.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the activation.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra kwargs for the normalization.
bias (bool, default: True ) –

Whether the projections have a bias term.

SparseBasicBlock [source] ¶

SparseBasicBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    bias: bool = False,
    indice_key: Optional[str] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SparseModule

Residual block of two submanifold sparse convolutions: conv, norm, act, conv, norm, add the input, act.

A pointwise submanifold convolution projects the skip connection when the channel count changes.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
kernel_size (int, default: 3 ) –

Kernel size of the two submanifold convolutions.
bias (bool, default: False ) –

Whether the two submanifold convolutions carry a bias.
indice_key (Optional[str], default: None ) –

spconv index key shared by the two submanifold convolutions.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

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

Extra keyword arguments for the normalization.

SparseConvBlock [source] ¶

SparseConvBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: Union[int, Tuple[int, ...]],
    *,
    stride: Union[int, Tuple[int, ...]] = 1,
    padding: Union[int, Tuple[int, ...]] = 0,
    indice_key: str,
    conv_type: str = "subm",
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SparseSequential

Sparse 3D convolution followed by normalization and activation.

A subm (submanifold), spconv (regular, optionally strided), or inverseconv convolution over a SparseConvTensor, followed by normalization and activation on the features. As a SparseSequential subclass it drops directly into the sparse voxel backbones' convolution stacks.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
kernel_size (Union[int, Tuple[int, ...]]) –

Convolution kernel size.
stride (Union[int, Tuple[int, ...]], default: 1 ) –

Convolution stride (used by spconv).
padding (Union[int, Tuple[int, ...]], default: 0 ) –

Convolution padding (used by spconv).
indice_key (str) –

spconv index key for the convolution.
conv_type (str, default: 'subm' ) –

One of subm, spconv, or inverseconv.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

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

Extra keyword arguments for the normalization (e.g. eps, momentum).
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.

SparseResidualBlock [source] ¶

SparseResidualBlock(
    in_channels: int,
    out_channels: int,
    *,
    indice_key: Optional[str] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SparseModule

Pre-activation sparse residual block shared by the SPFormer and SphereFormer SpConv U-Nets.

The convolutional branch applies \(\text{norm} \to \text{act} \to \text{conv} \to \text{norm} \to \text{act} \to \text{conv}\) over two \(3\times3\times3\) submanifold convolutions; the identity branch is nn.Identity when the channels match, else a \(1\times1\times1\) SubMConv3d projection.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
indice_key (Optional[str], default: None ) –

spconv index key shared by the two submanifold convolutions.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

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

Extra keyword arguments for the normalization (e.g. eps, momentum).

SubMConv3dBlock [source] ¶

SubMConv3dBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    padding: int,
    norm: Union[str, Callable, None] = None,
    act: Union[str, Callable, None] = None,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    stem_indice_key: Optional[str] = None,
)

Bases: Module

Submanifold sparse 3D convolution followed by normalization and activation, on packed point features.

Unlike SparseConvBlock, this block takes and returns packed \((N, C)\) features: it builds the SparseConvTensor from pos and batch itself, so it drops into a point-based backbone as a conditional position embedding.

Parameters:

in_channels (int) –

Number of input channels.
out_channels (int) –

Number of output channels.
kernel_size (int) –

Convolution kernel size.
padding (int) –

Convolution padding.
norm (Union[str, Callable, None], default: None ) –

Normalization passed to create_norm.
act (Union[str, Callable, None], default: None ) –

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

Extra keyword arguments for the activation.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

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

Whether the convolution has a bias term.
stem_indice_key (Optional[str], default: None ) –

spconv index key for the convolution.

SubMConv3dResidualBlock [source] ¶

SubMConv3dResidualBlock(
    channels: int,
    *,
    indice_key: str,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SparseModule

Residual block with a single \(3\times3\times3\) submanifold convolution: conv, norm, add the input, act.

Parameters:

channels (int) –

Input and output channels.
indice_key (str) –

Shared submanifold indice key.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

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

Extra normalization arguments.

DynamicTNet [source] ¶

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",
    track_running_stats: bool = True,
)

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 an EdgeConv, matching the Transform_Net from antao97/dgcnn.pytorch.

Architecture

edge_conv: EdgeConv 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.
track_running_stats (bool, default: True ) –

Whether to keep the per-sample transforms of the last forward pass (running_transform, \((B, k, k)\)) for losses.tnet_orthogonality_regularizer.

Methods:

forward –

Args:

forward [source] ¶

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

TNet [source] ¶

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",
    track_running_stats: bool = True,
)

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.
track_running_stats (bool, default: True ) –

Whether to keep the per-sample transforms of the last forward pass (running_transform, \((B, k, k)\)) for losses.tnet_orthogonality_regularizer.

Methods:

forward –

Forward pass of the T-Net.

forward [source] ¶

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

TorchSparseConvBlock [source] ¶

TorchSparseConvBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    transposed: bool = False,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Sparse 3D convolution followed by normalization and activation.

Set transposed=True for an upsampling (inverse) convolution.

TorchSparseResidualBlock [source] ¶

TorchSparseResidualBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.0,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Residual block of two sparse 3D convolutions.

A pointwise convolution projects the skip connection when the channel count or stride changes.

Attention [source] ¶

Attention(
    dim: int,
    num_heads: int = 8,
    qkv_bias: bool = False,
    qk_scale: Optional[float] = None,
    attn_dropout: float = 0.0,
    proj_dropout: float = 0.0,
)

Bases: Module

Multi-head self-attention over a dense token sequence.

Computes scaled dot-product attention with a single fused qkv projection and an output proj. An optional additive mask is added to the pre-softmax attention logits, which supports local / windowed attention (masked positions get a large negative bias).

Parameters:

dim (int) –

Token dimension \(C\). Must be divisible by num_heads.
num_heads (int, default: 8 ) –

Number of attention heads \(h\).
qkv_bias (bool, default: False ) –

Whether the fused query/key/value projection uses a bias.
qk_scale (Optional[float], default: None ) –

Override for the \(1/\sqrt{d_\text{head}}\) logit scale. Defaults to \(d_\text{head}^{-1/2}\) when None.
attn_dropout (float, default: 0.0 ) –

Dropout applied to the attention weights.
proj_dropout (float, default: 0.0 ) –

Dropout applied to the output projection.

Shape

Input: \((B, N, C)\) tokens and an optional mask broadcastable to \((B, h, N, N)\).
Output: \((B, N, C)\).

Example

import torch
from torch_pointcloud.layers import Attention

attn = Attention(384, num_heads=6)
x = torch.randn(2, 64, 384)
y = attn(x)
print(y.shape)

TransformerBlock [source] ¶

TransformerBlock(
    dim: int,
    num_heads: int,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = False,
    qk_scale: Optional[float] = None,
    dropout: float = 0.0,
    attn_dropout: float = 0.0,
    drop_path: float = 0.0,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = LayerNorm,
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Pre-norm transformer block: residual multi-head attention then a residual MLP.

Applies \(x \leftarrow x + \text{DropPath}(\text{Attn}(\text{Norm}(x)))\) followed by \(x \leftarrow x + \text{DropPath}(\text{MLP}(\text{Norm}(x)))\). The feed-forward is a plain-last torch_geometric.nn.MLP of hidden size \(\lfloor C \cdot \text{mlp\_ratio} \rfloor\), so activation and dropout are configurable through the resolver API.

Parameters:

dim (int) –

Token dimension \(C\).
num_heads (int) –

Number of attention heads.
mlp_ratio (float, default: 4.0 ) –

Hidden-to-input ratio of the feed-forward MLP.
qkv_bias (bool, default: False ) –

Whether the attention qkv projection uses a bias.
qk_scale (Optional[float], default: None ) –

Override for the attention logit scale.
dropout (float, default: 0.0 ) –

Dropout used in the MLP and the attention output projection.
attn_dropout (float, default: 0.0 ) –

Dropout applied to the attention weights.
drop_path (float, default: 0.0 ) –

Stochastic-depth rate for the two residual branches.
act (Union[str, Callable, None], default: 'gelu' ) –

Activation for the feed-forward MLP.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the activation.
norm (Union[str, Callable, None], default: LayerNorm ) –

Normalization applied before attention and before the MLP.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the normalization.

Shape

Input: \((B, N, C)\) tokens and an optional mask broadcastable to \((B, h, N, N)\).
Output: \((B, N, C)\).

Example

import torch
from torch_pointcloud.layers import TransformerBlock

block = TransformerBlock(384, num_heads=6, drop_path=0.1)
x = torch.randn(2, 64, 384)
y = block(x)
print(y.shape)

TransformerDecoderLayer [source] ¶

TransformerDecoderLayer(
    embed_dim: int,
    num_heads: int,
    mlp_dim: int,
    dropout: float,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_first: bool = True,
    pos_in_value: bool = False,
)

Bases: Module

DETR-style transformer decoder layer: self-attention over the queries, cross-attention onto the memory and a feed-forward block, each wrapped in a residual.

Query positions are added to the self-attention query and key and to the cross-attention query; memory positions are added to the cross-attention key. With pos_in_value both are added to the attention values as well. Tokens are sequence-first, \((S, B, C)\). With norm_first the layer norms are applied to the input of each sub-block (pre-norm); otherwise after each residual sum (post-norm).

Parameters:

embed_dim (int) –

Token embedding dimension.
num_heads (int) –

Number of attention heads.
mlp_dim (int) –

Hidden width of the feed-forward block.
dropout (float) –

Dropout probability.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the feed-forward block.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm_first (bool, default: True ) –

Whether to normalize the input of each sub-block rather than its residual sum.
pos_in_value (bool, default: False ) –

Whether the positional embeddings are also added to the attention values.

Shape

Input: \((T, B, C)\) queries, \((S, B, C)\) memory, \((S, B, C)\) memory positions, \((T, B, C)\) query positions and an optional \((B, S)\) memory key padding mask Output: \((T, B, C)\) queries

TransformerEncoderLayer [source] ¶

TransformerEncoderLayer(
    embed_dim: int,
    num_heads: int,
    mlp_dim: int,
    dropout: float,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_first: bool = True,
)

Bases: Module

DETR-style transformer encoder layer: self-attention with positional embeddings added to the query and key, then a feed-forward block, each wrapped in a residual.

Tokens are sequence-first, \((S, B, C)\), as nn.MultiheadAttention expects. With norm_first the layer norms are applied to the input of each sub-block (pre-norm); otherwise after each residual sum (post-norm).

Parameters:

embed_dim (int) –

Token embedding dimension.
num_heads (int) –

Number of attention heads.
mlp_dim (int) –

Hidden width of the feed-forward block.
dropout (float) –

Dropout probability.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the feed-forward block.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm_first (bool, default: True ) –

Whether to normalize the input of each sub-block rather than its residual sum.

Shape

Input: \((S, B, C)\) tokens, an optional \((S, S)\) attention mask, \((B, S)\) key padding mask and \((S, B, C)\) positions Output: \((S, B, C)\) tokens

DynamicMeanVFE [source] ¶

DynamicMeanVFE(
    in_channels: int,
    num_filters: Sequence[int],
    voxel_size: Sequence[float],
    point_cloud_range: Sequence[float],
    grid_size: Sequence[int],
)

Bases: Module

Dynamic mean voxel feature encoder for the Voxel Mamba detector.

Points are assigned to voxels on the fly (no fixed points-per-voxel), augmented with the per-voxel cluster-mean offset and voxel-center offset, then encoded by a stack of linear + norm + ReLU PFN layers whose per-voxel max-pool produces one feature vector per voxel.

Reference implementation: gwenzhang/Voxel-Mamba (DynamicVoxelVFE).

Parameters:

in_channels (int) –

Raw point feature channels including xyz (e.g. \(5\) for Waymo \(x, y, z, \text{intensity}, \text{elongation}\)).
num_filters (Sequence[int]) –

Output width of each PFN layer; the last entry is the voxel feature dim.
voxel_size (Sequence[float]) –

Voxel size \((v_x, v_y, v_z)\).
point_cloud_range (Sequence[float]) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
grid_size (Sequence[int]) –

Voxel grid extent \((n_x, n_y, n_z)\).

Shape

Input: \((N, C_\text{in})\) point features and \((N,)\) batch index.
Output: \((M, C_\text{out})\) voxel features and \((M, 4)\) voxel coords.

PillarFeatureLayer [source] ¶

PillarFeatureLayer(
    in_channels: int,
    out_channels: int,
    last: bool,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Pillar feature-net layer: linear + norm + act on every point, then a max-pool over each pillar.

Non-final layers halve their output width and concatenate the pooled feature back onto every point of the pillar; the final layer returns the pooled per-pillar feature. Padded pillars come as a dense \((P, N, C)\) tensor pooled over the point axis; dynamically voxelized points come packed as \((N, C)\) with their pillar_index, and are pooled by index.

Parameters:

in_channels (int) –

Input feature channels.
out_channels (int) –

Target output channels (halved internally for non-final layers).
last (bool) –

Whether this is the final layer.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

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

Extra normalization arguments.

Shape

Input: \((P, N, C_\text{in})\) padded pillars, or \((N, C_\text{in})\) packed points with an \((N,)\) pillar index.
Output: the input layout with \(C_\text{out}\) channels for non-final layers; \((P, 1, C_\text{out})\) or \((P, C_\text{out})\) pooled features for the final layer.

View [source] ¶

View(*shape: int)

Bases: Module

Returns a new view of the tensor with the same data. This is a wrapper around the torch.view function.

Note

This will not create a copy of the tensor and is generally more efficient than a torch.reshape call. However, it requires the input tensor to have a contiguous memory layout.

Parameters:

*shape (int, default: () ) –

The new shape of the tensor.

VoxelGridPool [source] ¶

VoxelGridPool(
    grid_size: float,
    reduce: str = "max",
    origin: Literal["grid", "min"] = "grid",
)

Bases: Module

Voxel grid subsampling: the points falling in the same grid_size cell are reduced to a single point.

Positions are averaged within a cell and features are reduced with reduce. The grid origin is either aligned to multiples of grid_size (origin="grid", so cell boundaries do not depend on the cloud's extent) or placed at each cloud's minimum corner (origin="min"). Passing return_inverse=True also returns the point-to-cell map.

Parameters:

grid_size (float) –

Cell edge length.
reduce (str, default: 'max' ) –

Feature reduction over a cell: "max", "mean", "sum" or "min".
origin (Literal['grid', 'min'], default: 'grid' ) –

Where the grid starts, "grid" or "min".

Shape

Input: \((N, C)\) features, \((N, 3)\) positions and \((N,)\) batch index.
Output: \((M, C)\) features, \((M, 3)\) positions, \((M,)\) batch index and, with return_inverse, the \((N,)\) cell index of every input point.

Example

>>> import torch
>>> from torch_pointcloud.layers import VoxelGridPool
>>> pool = VoxelGridPool(grid_size=0.5)
>>> x, pos, batch = torch.randn(100, 8), torch.rand(100, 3), torch.zeros(100, dtype=torch.long)
>>> x_pooled, pos_pooled, batch_pooled = pool(x, pos, batch)  # doctest: +SKIP
>>> x_pooled.shape[1], pos_pooled.shape[1]  # doctest: +SKIP
(8, 3)

XConv [source] ¶

XConv(
    in_channels: int,
    out_channels: int,
    spatial_dim: int,
    kernel_size: int,
    hidden_channels: Optional[int] = None,
    depth_multiplier: Optional[int] = None,
    dilation: int = 1,
    act: Union[str, Callable, None] = "elu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
)

Bases: Module

XConv layer as described in the paper "PointCNN: Convolution On X-Transformed Points" by Yangyan Li, Rui Bu, Mingchao Sun, Wei Wu, Xinhan Di, Baoquan Chen.

This layer is inspired by the PyTorch Geometric torch_geometric.nn.XConv layer implementation, excepts that this convolution layer supports bipartite graphs and more flexibility / customization. This implementation provides full compatibility with the PyTorch Geometric library while following the official paper and original tensorflow implementation.

edge_index lists [source, target] pairs with exactly kernel_size * dilation consecutive edges per target point, in the order knn returns them; a bipartite graph is given as (x_source, None) and (pos_source, pos_target) pairs.

Methods:

init_weights_ –

Applies Xavier normal initialization to nn.Linear and nn.Conv1d modules, zeroing their bias.

init_weights_ `staticmethod` [source] ¶

init_weights_(module: Module) -> None

Applies Xavier normal initialization to nn.Linear and nn.Conv1d modules, zeroing their bias.

create_module [source] ¶

create_module(
    name: ModuleLike,
    *args: Any,
    registry: ModuleRegistryDict,
    **kwargs: Any,
) -> Module

Utility function to instantiate modules from a registry, or directly pass a module instance.

Parameters:

name (ModuleLike) –

The name of the module to instantiate, or the module type or instance itself.
args (Any, default: () ) –

Positional arguments to pass to the module constructor.
registry (ModuleRegistryDict) –

A dictionary mapping module names to their types.
kwargs (Any, default: {} ) –

Keyword arguments to pass to the module constructor.

Returns:

Module –

The instantiated module.

create_act [source] ¶

create_act(
    act: Union[str, Callable, None], **act_kwargs: Any
) -> Optional[Module]

Resolve an activation, returning None for act=None.

Parameters:

act (Union[str, Callable, None]) –

Activation name (e.g., "relu", "leaky_relu"; "swish" resolves to nn.SiLU), a class, an instance, or None.
**act_kwargs (Any, default: {} ) –

Forwarded to the activation constructor (ignored if act is already an instance).

Returns:

Optional[Module] –

The instantiated activation module, or None if act is None.

prediction_branch [source] ¶

prediction_branch(
    in_channels: int,
    out_channels: int,
    num_middle_conv: int,
    num_middle_filter: int,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
) -> Sequential

Build a prediction branch: middle conv blocks, then a plain output conv.

The per-attribute branch of the multi-head detectors: num_middle_conv blocks of (\(3\times3\) conv, norm, act) followed by a \(3\times3\) output conv.

Parameters:

in_channels (int) –

Input channels.
out_channels (int) –

Output channels of the final conv.
num_middle_conv (int) –

Number of middle conv blocks.
num_middle_filter (int) –

Channel width of the middle convs.
act (Union[str, Callable, None], default: 'relu' ) –

Activation of the middle conv blocks.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization of the middle conv blocks.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.
bias (bool, default: False ) –

Whether the middle convs carry a bias (the output conv always does).

Returns:

Sequential –

The branch as an nn.Sequential.

Shape

Input: \((B, C_\text{in}, H, W)\)
Output: \((B, C_\text{out}, H, W)\)

Example

>>> branch = prediction_branch(64, 2, num_middle_conv=1, num_middle_filter=64)
>>> branch(torch.rand(2, 64, 16, 16)).shape
torch.Size([2, 2, 16, 16])

create_dropout [source] ¶

create_dropout(
    name: None, *args: Any, **kwargs: Any
) -> None

create_dropout(
    name: Union[DropoutLike, float],
    *args: Any,
    **kwargs: Any,
) -> Module

create_dropout(
    name: Union[DropoutLike, float, None],
    *args: Any,
    **kwargs: Any,
) -> Optional[Module]

Resolve a dropout layer from a name, a class, an instance, or a probability, returning None for name=None.

Parameters:

name (Union[DropoutLike, float, None]) –

Dropout name (e.g., "dropout", "drop_path"), a class, an instance, a float probability, which is a shorthand for nn.Dropout(p=name), or None.
*args (Any, default: () ) –

Forwarded to the dropout constructor.
**kwargs (Any, default: {} ) –

Forwarded to the dropout constructor (ignored if name is already an instance).

Returns:

Optional[Module] –

The instantiated dropout module, or None if name is None.

drop_path [source] ¶

drop_path(
    x: Tensor,
    drop_prob: float = 0.0,
    training: bool = False,
    scale_by_keep: bool = True,
) -> Tensor

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks), as described in the paper Deep Networks with Stochastic Depth by Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, Kilian Weinberger.

Implementation is taken from original implementation by Ross Wightman in pytorch-image-models.

create_kernel_points [source] ¶

create_kernel_points(
    radius: float,
    num_points: int,
    fixed_position: Literal[
        "none", "center", "vertical"
    ] = "center",
    method: Literal["lloyd", "gradient"] = "lloyd",
) -> Tensor

Builds the kernel point positions of a KPConv kernel, randomly rotated and jittered.

Positions are optimized on the unit sphere, cached under CACHE_DIR and reused across calls.

Parameters:

radius (float) –

Radius of the sphere the kernel points are scaled to.
num_points (int) –

Number of kernel points \(K\).
fixed_position (Literal['none', 'center', 'vertical'], default: 'center' ) –

Which kernel point is pinned: "none", "center", or "vertical".
method (Literal['lloyd', 'gradient'], default: 'lloyd' ) –

Optimization used to spread the points, either "lloyd" or "gradient".

Returns:

Tensor –

Kernel point positions of shape \((K, 3)\).

create_norm [source] ¶

create_norm(
    norm: Union[str, Callable, None],
    channels: int,
    *,
    dim: int = 1,
    conditions: Optional[Sequence[str]] = None,
    **norm_kwargs: Any,
) -> Optional[Module]

Resolve a normalization layer with a spatial dimensionality hint.

For dim == 1, defers to PyG's normalization_resolver (graph-aware norms plus BatchNorm1d). For dim \(\in \{2, 3\}\), maps common norm names to the matching nn.*Nd variant. Pass a class or an existing instance to bypass string resolution. When conditions is given, the resolved norm is wrapped in a per-condition PDNorm for multi-dataset (prompt-driven) training.

Parameters:

norm (Union[str, Callable, None]) –

Norm name ("batch_norm", "instance_norm", "group_norm", "layer_norm"), a class, an instance, or None.
channels (int) –

Number of feature channels.
dim (int, default: 1 ) –

Spatial dimensionality of the feature map. \(1\) for packed / graph tensors \((N, C)\), \(2\) for \((B, C, H, W)\), \(3\) for \((B, C, H, W, D)\).
conditions (Optional[Sequence[str]], default: None ) –

Ordered dataset condition names. When set (and norm is not None), returns a PDNorm holding one inner norm per condition.
**norm_kwargs (Any, default: {} ) –

Forwarded to the norm constructor.

Returns:

Optional[Module] –

The instantiated norm module, or None if norm is None.

create_adaptive_pool [source] ¶

create_adaptive_pool(
    name: None, *args: Any, **kwargs: Any
) -> None

create_adaptive_pool(
    name: AdaptivePoolLike, *args: Any, **kwargs: Any
) -> Module

create_adaptive_pool(
    name: Optional[AdaptivePoolLike],
    *args: Any,
    **kwargs: Any,
) -> Optional[Module]

Resolve a dense adaptive pooling module (nn.AdaptiveAvgPool1d / nn.AdaptiveMaxPool1d), returning None for name=None.

Parameters:

name (Optional[AdaptivePoolLike]) –

Pool name ("mean", "max"), a module class, an existing instance (returned as-is), or None.
*args (Any, default: () ) –

Positional arguments forwarded to the pool constructor.
**kwargs (Any, default: {} ) –

Keyword arguments forwarded to the pool constructor. output_size defaults to 1 (global pooling).

Returns:

Optional[Module] –

The instantiated pooling module, or None if name is None.

create_pool [source] ¶

create_pool(name: None, *args: Any, **kwargs: Any) -> None

create_pool(
    name: PoolLike, *args: Any, **kwargs: Any
) -> Module

create_pool(
    name: Optional[PoolLike], *args: Any, **kwargs: Any
) -> Optional[Module]

Resolve a packed-batch pooling module from a name, class, or instance, returning None for name=None.

Parameters:

name (Optional[PoolLike]) –

Pool name ("max", "min", "mean", "mul", "sum"), a module class, an existing instance (returned as-is), or None.
*args (Any, default: () ) –

Positional arguments forwarded to the pool constructor.
**kwargs (Any, default: {} ) –

Keyword arguments forwarded to the pool constructor (dim, dim_size).

Returns:

Optional[Module] –

The instantiated pooling module, or None if name is None.

patch_layout [source] ¶

patch_layout(
    batch: Tensor, patch_size: int, mode: PadMode
) -> PatchLayout

Pad every sample to whole patches of patch_size points and derive the patch offsets.

Flash attention handles a sample smaller than the patch as a short variable-length patch, so it only pads samples at or above the patch size (mode="above"); the dense reshape needs every sample padded to a full patch (mode="all"). Padding, never shrinking, keeps one sample's output independent of its co-batched neighbors.

Parameters:

batch (Tensor) –

Batch index of every point, sorted, shape \((N,)\).
patch_size (int) –

Number of points per attention patch.
mode (PadMode) –

Which samples to pad, see divisible_pad.

Returns:

PatchLayout –

The PatchLayout of the points.

Overview

Affine [source] ¶

AnchorGroupHead [source] ¶

reset_parameters [source] ¶

AnchorHead [source] ¶

reset_parameters [source] ¶

generate_predicted_boxes [source] ¶

decode [source] ¶

AnchorHeadMultiOutput [source] ¶

AnchorHeadOutput [source] ¶

MultiGroupAnchorHead [source] ¶

decode [source] ¶

BEVBackbone [source] ¶

BEVResidualBackbone [source] ¶

ResidualBlock2d [source] ¶

Conv2dBlock [source] ¶

Conv3dBlock [source] ¶

DropPath [source] ¶

FPS [source] ¶

GeometricAffineConv [source] ¶

GridPool [source] ¶

forward [source] ¶

GroupedVectorAttention [source] ¶

KPConv [source] ¶

deformable property ¶

configure_offsets [source] ¶

configure_kernel [source] ¶

KPConvBlock [source] ¶

KPResidualBlock [source] ¶

LinearBlock [source] ¶

OctreeAttention [source] ¶

dilated property ¶

forward_attn [source] ¶

forward_proj [source] ¶

forward_rpe [source] ¶

OctreeRelativePositionEncoding [source] ¶

pos_to_idx [source] ¶

OctreeT [source] ¶

block_size property ¶

reset [source] ¶

from_octree classmethod [source] ¶

construct_all_attention_context [source] ¶

construct_attention_context [source] ¶

pad_to_patch_size [source] ¶

unpad [source] ¶

OctreeConvBlock [source] ¶

OctreeDeconvBlock [source] ¶

PDNorm [source] ¶

PointPatchEmbed [source] ¶

PointTransformerConv [source] ¶

PointConv [source] ¶

PointConvDensity [source] ¶

PointConvDensityGlobalSetAbstraction [source] ¶

PointConvDensitySetAbstraction [source] ¶

PointConvGlobalSetAbstraction [source] ¶

PointConvSetAbstraction [source] ¶

PointNet2Conv [source] ¶

PointNet2FeaturePropagation [source] ¶

PointNet2GlobalSetAbstraction [source] ¶

PointNet2SetAbstraction [source] ¶

PointNeXtConv [source] ¶

PointNeXtResidualBlock [source] ¶

PointNeXtSetAbstraction [source] ¶

CatPool [source] ¶

num_pools property ¶

MaxPool [source] ¶

MeanPool [source] ¶

MinPool [source] ¶

MulPool [source] ¶

SumPool [source] ¶

PVConv [source] ¶

SE3d [source] ¶

Voxelization [source] ¶

AttentivePooling [source] ¶

LocalFeatureAggregation [source] ¶

LocalSpatialEncoding [source] ¶

RandLANetResidualBlock [source] ¶

Point3DRoPE [source] ¶

PatchLayout [source] ¶

RelativePositionalEncoding [source] ¶

deformable `property` ¶

dilated `property` ¶

block_size `property` ¶

from_octree `classmethod` [source] ¶

num_pools `property` ¶

init_weights_ `staticmethod` [source] ¶