Skip to content

Overview

Model architectures, the create_model factory, and the pretrained weight registry.

Modules:

  • concerto –

    Concerto pretrained encoders and linear-probing segmentation model.

  • dgcnn –

    DGCNN classification and segmentation models.

  • kpconv –

    KPConv classification and segmentation models.

  • lion –

    LION detection model.

  • octformer –

    OctFormer classification and segmentation models.

  • point_bert –

    Point-BERT classification models, discrete VAE tokenizer, and masked pretraining transformer.

  • point_m2ae –

    Point-M2AE classification, segmentation, and masked autoencoder pretraining models.

  • point_mae –

    Point-MAE classification, segmentation, and masked autoencoder pretraining models.

  • point_mamba –

    Point-Mamba classification and masked autoencoder pretraining models.

  • point_transformer –

    Point Transformer classification and segmentation models.

  • point_transformer_v2 –

    Point Transformer V2 classification and segmentation models.

  • point_transformer_v3 –

    Point Transformer V3 classification and segmentation models.

  • pointcnn –

    PointCNN classification and segmentation models.

  • pointconv –

    PointConv classification model.

  • pointgpt –

    PointGPT classification and generative pretraining models.

  • pointmlp –

    PointMLP classification and segmentation models.

  • pointnet –

    PointNet classification and segmentation models.

  • pointnet2 –

    PointNet++ classification and segmentation models.

  • pointnext –

    PointNeXt classification and segmentation models.

  • pointpillars –

    PointPillars detection models.

  • pointrcnn –

    PointRCNN detection model.

  • pvcnn –

    PVCNN classification and segmentation models.

  • pvcnn2 –

    PVCNN++ classification and segmentation models.

  • randlanet –

    RandLA-Net classification and segmentation models.

  • second –

    SECOND detection models.

  • sonata –

    Sonata pretrained encoder and linear-probing segmentation model.

  • spformer_unet –

    SPFormer-UNet segmentation model.

  • sphereformer –

    SphereFormer segmentation model.

  • spunet –

    SpUNet segmentation model.

  • spvcnn –

    SPVCNN classification and segmentation models.

  • threedetr –

    3DETR: an end-to-end transformer detector for 3D point clouds.

  • utonia –

    Utonia pretrained encoder and linear-probing segmentation model.

  • votenet –

    VoteNet detection model.

  • voxel_mamba –

    Voxel-Mamba detection model.

  • voxelnext –

    VoxelNeXt detection model.

Classes:

Functions:

  • create_model –

    Build a registered model, optionally loading pretrained or local checkpoint weights.

  • deprecated_model –

    Keep a renamed model reachable under its former name, placed above its register_model decorator.

  • list_models –

    List registered model names, sorted alphabetically.

  • model_info –

    The registry entry of a model without building it: its task, weights, transform, input keys, hparams and

  • register_model –

    Register a model entry point under name for task.

ClassificationModel [source] ΒΆ

ClassificationModel(in_channels: int, num_classes: int)

Bases: Module

Base class for point cloud classification models.

Subclasses implement forward_features (encode the packed cloud), forward_head (pool and map to logits), configure_head and reset_classifier, and compose the first two in forward. Building a model with num_classes=0 makes configure_head return nn.Identity, so forward returns the pooled features \((B, \text{num\_features})\); forward_head(..., pre_logits=True) returns those same features for any num_classes.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of output classes; 0 drops the head.

Methods:

  • configure_head –

    Build and return the classification head for the current num_classes (nn.Identity when 0).

  • reset_classifier –

    Replace the classification head for num_classes outputs (0 drops it).

  • forward –

    Run the model on a packed point cloud.

  • forward_features –

    Encode a packed point cloud into features for forward_head.

  • forward_head –

    Map encoded features to class logits.

Attributes:

  • num_features (int) –

    Width \(C\) of the pooled features entering the head.

num_features abstractmethod property ΒΆ

num_features: int

Width \(C\) of the pooled features entering the head.

The last dimension of forward(...) when num_classes=0 and of forward_head(..., pre_logits=True). Read-only: it is derived from the instantiated submodules.

configure_head abstractmethod [source] ΒΆ

configure_head() -> Module

Build and return the classification head for the current num_classes (nn.Identity when 0).

Called from both __init__ and reset_classifier so the two always build the same module.

reset_classifier abstractmethod [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification head for num_classes outputs (0 drops it).

Concrete models may accept extra keyword options (e.g. global_pool, head_channels, dropout).

forward abstractmethod [source] ΒΆ

forward(*args: Any, **kwargs: Any) -> Any

Run the model on a packed point cloud.

Canonical signature: forward(x, pos, batch) with features x \((N, C)\) (or None), positions pos \((N, 3)\) and batch index batch \((N,)\), returning logits \((B, \text{num\_classes})\), or the pooled features \((B, \text{num\_features})\) when num_classes=0.

forward_features abstractmethod [source] ΒΆ

forward_features(*args: Any, **kwargs: Any) -> Any

Encode a packed point cloud into features for forward_head.

Canonical signature: forward_features(x, pos, batch, return_intermediates=False), returning the encoded features (often with the downsampled pos / batch). With return_intermediates=True, the finer encoder stages are appended as the last element of the return: a List[FeaturesDict] in fine-to-coarse order, each entry holding x, batch and the stage's pos or pos_grid. Models add the keys their decoder needs or use their own container when a stage is not a packed point set (dense tokens, voxel-point pairs).

forward_head abstractmethod [source] ΒΆ

forward_head(*args: Any, **kwargs: Any) -> Any

Map encoded features to class logits.

Canonical signature: forward_head(x, batch, pre_logits=False), returning logits \((B, \text{num\_classes})\), or the pooled features \((B, \text{num\_features})\) when pre_logits=True.

DetectionModel [source] ΒΆ

DetectionModel(in_channels: int, num_classes: int)

Bases: Module

Base class for 3D object detection models.

Subclasses implement forward_features (the backbone), forward_head (backbone features to the raw prediction output), decode (raw output to packed boxes) and num_features, and compose the first two in forward. decode is raw: score thresholding and NMS belong to the evaluation postprocess, not the model. Detectors have no headless mode: feature maps come from forward_features, and reset_classifier is only available where the class count can change without touching the box branches.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of object classes.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • forward –

    Run the model and return its raw prediction output.

  • forward_features –

    Run the backbone on the packed input.

  • forward_head –

    Map backbone features to the raw prediction output.

  • decode –

    Decode the raw forward output into packed boxes.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the forward_features output.

num_features abstractmethod property ΒΆ

num_features: int

Channel count \(C\) of the forward_features output.

The BEV map channels, the sparse voxel feature width, or the packed point feature width, depending on the backbone. Read-only: it is derived from the instantiated submodules.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

forward abstractmethod [source] ΒΆ

forward(*args: Any, **kwargs: Any) -> Any

Run the model and return its raw prediction output.

Canonical signature: forward(x, pos, batch) for point-based detectors (voxel detectors take their voxelized layout instead), returning the prediction dict consumed by the training criterion and by decode.

forward_features abstractmethod [source] ΒΆ

forward_features(*args: Any, **kwargs: Any) -> Any

Run the backbone on the packed input.

Canonical signature: forward_features(x, pos, batch) for point-based detectors (voxel detectors take their voxelized layout instead), returning the backbone features fed to the detection head.

forward_head abstractmethod [source] ΒΆ

forward_head(*args: Any, **kwargs: Any) -> Any

Map backbone features to the raw prediction output.

Canonical signature: forward_head(features), returning the same prediction output as forward. Heads that need more than the backbone output take it as extra arguments (VoteNet seed indices, 3DETR scene extents, PointRCNN ground truth in train mode).

decode abstractmethod [source] ΒΆ

decode(*args: Any, **kwargs: Any) -> Any

Decode the raw forward output into packed boxes.

Canonical signature: decode(output), returning a Detection3D with boxes \((K, 7)\), scores \((K,)\), labels \((K,)\) and batch \((K,)\). The decode is raw: no score thresholding and no NMS.

PartSegmentationModel [source] ΒΆ

PartSegmentationModel(
    in_channels: int, num_classes: int, num_categories: int
)

Bases: Module

Base class for part segmentation models, registered under the "part-segmentation" task.

Part segmentation labels the parts of a single object and conditions the prediction on the object category, passed to forward as a one-hot \((B, \text{num\_categories})\).

Subclasses implement forward_features (encoder, keeping the skip intermediates), forward_head (per-point logits), configure_head and reset_classifier, and compose them in forward, usually through forward_decoder (upsampling path). Building a model with num_classes=0 makes configure_head return nn.Identity, so forward returns the per-point features \((N, \text{num\_features})\); forward_head(..., pre_logits=True) returns those same features for any num_classes.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of part classes; 0 drops the head.

  • num_categories (int) –

    Number of object categories of the one-hot category input.

Methods:

  • configure_head –

    Build and return the segmentation head for the current num_classes (nn.Identity when 0).

  • reset_classifier –

    Replace the segmentation head for num_classes outputs (0 drops it).

  • forward –

    Run the model on a packed point cloud.

  • forward_features –

    Encode a packed point cloud, keeping what forward_decoder needs for the skip connections.

  • forward_decoder –

    Decode encoder features back to per-point resolution.

  • forward_head –

    Map per-point features to per-point logits.

Attributes:

  • num_features (int) –

    Width \(C\) of the per-point features entering the head.

num_features abstractmethod property ΒΆ

num_features: int

Width \(C\) of the per-point features entering the head.

The last dimension of forward(...) when num_classes=0 and of forward_head(..., pre_logits=True). Read-only: it is derived from the instantiated submodules.

configure_head abstractmethod [source] ΒΆ

configure_head() -> Module

Build and return the segmentation head for the current num_classes (nn.Identity when 0).

Called from both __init__ and reset_classifier so the two always build the same module.

reset_classifier abstractmethod [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the segmentation head for num_classes outputs (0 drops it).

Concrete models may accept extra keyword options (e.g. global_pool, head_channels, dropout).

forward abstractmethod [source] ΒΆ

forward(*args: Any, **kwargs: Any) -> Any

Run the model on a packed point cloud.

Canonical signature: forward(x, pos, batch, category) with features x \((N, C)\) (or None), positions pos \((N, 3)\), batch index batch \((N,)\) and the one-hot object category \((B, \text{num\_categories})\), returning per-point logits \((N, \text{num\_classes})\), or the per-point features \((N, \text{num\_features})\) when num_classes=0.

forward_features abstractmethod [source] ΒΆ

forward_features(*args: Any, **kwargs: Any) -> Any

Encode a packed point cloud, keeping what forward_decoder needs for the skip connections.

Canonical signature: forward_features(x, pos, batch, return_intermediates=False), returning the encoder output. With return_intermediates=True, the finer encoder stages are appended as the last element of the return: a List[FeaturesDict] in fine-to-coarse order, each entry holding x, batch and the stage's pos or pos_grid. Models add the keys their decoder needs (e.g. pooling_inverse) or use their own container when a stage is not a packed point set (dense tokens, voxel-point pairs).

forward_decoder [source] ΒΆ

forward_decoder(*args: Any, **kwargs: Any) -> Any

Decode encoder features back to per-point resolution.

Canonical signature: forward_decoder(x, ..., intermediates), consuming the output of forward_features and returning per-point features \((N, C)\). Models whose encoder already emits per-point features (DGCNN, PointNet) have no decoder and raise NotImplementedError.

forward_head abstractmethod [source] ΒΆ

forward_head(*args: Any, **kwargs: Any) -> Any

Map per-point features to per-point logits.

Canonical signature: forward_head(x, pre_logits=False), returning logits \((N, \text{num\_classes})\), or the per-point features \((N, \text{num\_features})\) when pre_logits=True.

PretrainingModel [source] ΒΆ

PretrainingModel(in_channels: int)

Bases: Module

Base class for self-supervised pretraining models, registered under the "pretraining" task.

Pretraining models (masked autoencoders, generative pretraining) return pretext outputs from forward, such as (pred, target) group coordinates for a reconstruction objective, instead of logits or detections. The class only stores in_channels; each subclass defines its own forward signature.

Models with a task head inherit from one of the task ABCs instead: ClassificationModel, SemanticSegmentationModel, PartSegmentationModel or DetectionModel. The classification and segmentation ABCs require the split into forward_features (encoder), forward_decoder (segmentation only, optional) and forward_head (logits), plus configure_head / reset_classifier for building the head and the read-only num_features property for the width of the features entering it. Detection models split into forward_features (backbone), forward_head (raw predictions) and decode (boxes), and expose num_features as the backbone output width.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

SemanticSegmentationModel [source] ΒΆ

SemanticSegmentationModel(
    in_channels: int, num_classes: int
)

Bases: Module

Base class for point cloud semantic segmentation models, registered under the "semantic-segmentation" task.

Subclasses implement forward_features (encoder, keeping the skip intermediates), forward_head (per-point logits), configure_head and reset_classifier, and compose them in forward, usually through forward_decoder (upsampling path). Building a model with num_classes=0 makes configure_head return nn.Identity, so forward returns the per-point features \((N, \text{num\_features})\); forward_head(..., pre_logits=True) returns those same features for any num_classes.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of output classes; 0 drops the head.

Methods:

  • configure_head –

    Build and return the segmentation head for the current num_classes (nn.Identity when 0).

  • reset_classifier –

    Replace the segmentation head for num_classes outputs (0 drops it).

  • forward –

    Run the model on a packed point cloud.

  • forward_features –

    Encode a packed point cloud, keeping what forward_decoder needs for the skip connections.

  • forward_decoder –

    Decode encoder features back to per-point resolution.

  • forward_head –

    Map per-point features to per-point logits.

Attributes:

  • num_features (int) –

    Width \(C\) of the per-point features entering the head.

num_features abstractmethod property ΒΆ

num_features: int

Width \(C\) of the per-point features entering the head.

The last dimension of forward(...) when num_classes=0 and of forward_head(..., pre_logits=True). Read-only: it is derived from the instantiated submodules.

configure_head abstractmethod [source] ΒΆ

configure_head() -> Module

Build and return the segmentation head for the current num_classes (nn.Identity when 0).

Called from both __init__ and reset_classifier so the two always build the same module.

reset_classifier abstractmethod [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the segmentation head for num_classes outputs (0 drops it).

Concrete models may accept extra keyword options (e.g. global_pool, head_channels, dropout).

forward abstractmethod [source] ΒΆ

forward(*args: Any, **kwargs: Any) -> Any

Run the model on a packed point cloud.

Canonical signature: forward(x, pos, batch) with features x \((N, C)\) (or None), positions pos \((N, 3)\) and batch index batch \((N,)\), returning per-point logits \((N, \text{num\_classes})\), or the per-point features \((N, \text{num\_features})\) when num_classes=0.

forward_features abstractmethod [source] ΒΆ

forward_features(*args: Any, **kwargs: Any) -> Any

Encode a packed point cloud, keeping what forward_decoder needs for the skip connections.

Canonical signature: forward_features(x, pos, batch, return_intermediates=False), returning the encoder output. With return_intermediates=True, the finer encoder stages are appended as the last element of the return: a List[FeaturesDict] in fine-to-coarse order, each entry holding x, batch and the stage's pos or pos_grid. Models add the keys their decoder needs (e.g. pooling_inverse) or use their own container when a stage is not a packed point set (dense tokens, voxel-point pairs).

forward_decoder [source] ΒΆ

forward_decoder(*args: Any, **kwargs: Any) -> Any

Decode encoder features back to per-point resolution.

Canonical signature: forward_decoder(x, ..., intermediates), consuming the output of forward_features and returning per-point features \((N, C)\). Models whose encoder already emits per-point features (DGCNN, PointNet) have no decoder and raise NotImplementedError.

forward_head abstractmethod [source] ΒΆ

forward_head(*args: Any, **kwargs: Any) -> Any

Map per-point features to per-point logits.

Canonical signature: forward_head(x, pre_logits=False), returning logits \((N, \text{num\_classes})\), or the per-point features \((N, \text{num\_features})\) when pre_logits=True.

DeprecatedModelDict [source] ΒΆ

Bases: TypedDict

Former name of a registered model, declared with deprecated_model.

Attributes:

  • name (str) –

    The former registry name.

  • deprecated_in (str) –

    Release that renamed the model (e.g. 0.0.17).

  • removed_in (Optional[str]) –

    Release that drops the former name, or None when it is not scheduled.

  • message (Optional[str]) –

    Sentence appended to the warning the former name raises, or None.

ModelDict [source] ΒΆ

Bases: TypedDict

Registry entry of a registered model: its entry point and the metadata create_model returns.

ModelInfoDict [source] ΒΆ

Bases: TypedDict

Registry entry of a built model, returned by create_model(..., return_info=True).

WeightsDict [source] ΒΆ

Bases: TypedDict

Structured description of a pretrained checkpoint.

Attributes:

  • url (ValueCollection[str]) –

    Location of the weight file, mirroring the Hub path so that hf:// stands for https://huggingface.co/ (e.g. hf://torch-pointcloud/pointnext-sm.scanobjectnn-hardest.openpoints/resolve/a98ba0b8f7c9615acf0f38a1a20c97bfb01c2e79/model.safetensors). Several URLs are read in order into a single state dict (e.g. an encoder file and its head file). A file in a gated repository downloads with the Hub token of the environment (see load_hub_token).

  • pretrained_filter_fn (NotRequired[Callable[[Dict[str, Any], Module], Dict[str, Any]]]) –

    Called as pretrained_filter_fn(state_dict, model) on the state dict read from url before it is loaded, to rename, drop or reorder the entries of a checkpoint saved with another parameter layout (e.g. the authors' original release).

  • dataset (NotRequired[str]) –

    Benchmark the checkpoint was trained on (e.g. scanobjectnn, s3dis-area5).

  • metrics (NotRequired[Dict[str, float]]) –

    Scores measured with this package's benchmark scripts, keyed by metric name (e.g. {"OA": 88.20}, {"mIoU": 0.7604}).

  • classes (NotRequired[Sequence[str]]) –

    Label names in prediction-channel order, so classes[i] names class \(i\) of the head.

  • author (NotRequired[str]) –

    Tag identifying who trained the original checkpoint (e.g. openpcdet, fair).

  • license (NotRequired[str]) –

    License of the weight file (e.g. CC-BY-NC-4.0), when the source repository declares one.

ConcertoSegmentation [source] ΒΆ

ConcertoSegmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "z",
        "z-trans",
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (3, 3, 3, 12, 3),
    encoder_channels: Sequence[int] = (
        64,
        128,
        256,
        512,
        768,
    ),
    encoder_num_heads: Sequence[int] = (4, 8, 16, 32, 48),
    encoder_patch_size: Sequence[int] = (
        1024,
        1024,
        1024,
        1024,
        1024,
    ),
    norm: Union[str, Callable] = "layer_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    dropout: float = 0.0,
    pooling: str = "grid",
    stem_type: str = "linear",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
)

Bases: SemanticSegmentationModel

Concerto linear-probing segmentation model.

Linear-probe variant from Concerto: Joint 2D-3D Self-Supervised Learning (Zhang et al., NeurIPS 2025). The pretrained PT-V3 encoder is run with intermediate features captured at every encoder stage; those features are successively unpooled through the saved pooling inverses, concatenated with each parent stage, then projected by a single linear segmentation head.

Architecturally identical to Sonata's linear-probe head, only the encoder scale differs (Concerto-large uses \(C = (64, 128, 256, 512, 768)\) versus Sonata's \((48, 96, 192, 384, 512)\)).

Note

The default (and registered) configuration enables flash attention (use_flash_attn=True), which requires the flash-attn package and a CUDA device; pass use_flash_attn=False to run without it. The xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run the forward on CPU.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

DGCNNClassification [source] ΒΆ

DGCNNClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stnet_local_channels: Optional[Sequence[int]] = None,
    stnet_global_channels: Optional[Sequence[int]] = None,
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    channels: Sequence[int],
    proj_channels: Optional[int] = None,
    num_neighbors: Union[int, Sequence[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,
    global_pool: PoolLike | Sequence[PoolLike] = "max",
)

Bases: ClassificationModel

Classification model as described in the paper "Dynamic Graph CNN for Learning on Point Clouds" by Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, Justin M. Solomon.

DGCNN introduces the EdgeConv operator, which computes features on dynamically constructed k-nearest neighbor (k-NN) graphs at each layer. Graphs are recomputed in the learned feature space, allowing the network to capture both local geometric relationships and long-range semantic structures.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

  • stnet_local_channels (Optional[Sequence[int]], default: None ) –

    List of channels for the local spatial transformer network. If None, the spatial transformer network is not used.

  • stnet_global_channels (Optional[Sequence[int]], default: None ) –

    List of channels for the global spatial transformer network. If None, the spatial transformer network is not used.

  • channels (Sequence[int]) –

    List of channels for each encoder block.

  • proj_channels (Optional[int], default: None ) –

    If set, projects the concatenated encoder features through an MLP of this width before pooling. Matches the conv5 layer in the original DGCNN paper.

  • head_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    List of channels for each head block.

  • num_neighbors (Union[int, Sequence[int]]) –

    Maximum number of neighbors for each encoder block.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the activation function.

  • act_first (bool, default: False ) –

    Whether to apply activation before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the normalization layer.

  • bias (bool, default: True ) –

    Whether to use bias in linear / MLP layers.

  • dropout (float, default: 0.0 ) –

    Dropout probability before the classification head.

  • global_pool (PoolLike | Sequence[PoolLike], default: 'max' ) –

    Global pooling operation for final feature aggregation.

Methods:

  • configure_stnet –

    Build the spatial transformer aligning the input coordinates, or None when its channels are unset.

  • configure_encoder –

    Build the DGCNNEncoder backbone.

  • configure_proj –

    Build the projection MLP applied before pooling, or None when proj_channels is unset.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_stnet [source] ΒΆ

configure_stnet() -> Optional[TNet]

Build the spatial transformer aligning the input coordinates, or None when its channels are unset.

configure_encoder [source] ΒΆ

configure_encoder() -> DGCNNEncoder

Build the DGCNNEncoder backbone.

configure_proj [source] ΒΆ

configure_proj() -> Optional[MLP]

Build the projection MLP applied before pooling, or None when proj_channels is unset.

DGCNNPartSegmentation [source] ΒΆ

DGCNNPartSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    num_categories: int,
    cat_embed_channels: int = 64,
    spatial_dim: int = 3,
    stnet_edge_channels: Optional[Sequence[int]] = None,
    stnet_local_channels: Optional[Sequence[int]] = None,
    stnet_global_channels: Optional[Sequence[int]] = None,
    stnet_num_neighbors: int = 20,
    proj_channels: int = 1024,
    channels: Sequence[Union[int, Sequence[int]]],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    num_neighbors: Union[int, Sequence[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,
)

Bases: PartSegmentationModel

Part segmentation model as described in the paper "Dynamic Graph CNN for Learning on Point Clouds" by Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, Justin M. Solomon.

Extends the DGCNN encoder with a category-conditioned global feature branch for part-level segmentation (e.g. ShapeNet parts).

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output part classes (across all categories).

  • num_categories (int) –

    Number of object categories for the category embedding.

  • cat_embed_channels (int, default: 64 ) –

    Number of channels for the category embedding.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

  • stnet_edge_channels (Optional[Sequence[int]], default: None ) –

    Hidden channels for the DynamicTNet EdgeConv MLP. If None, the spatial transformer network is not used.

  • stnet_local_channels (Optional[Sequence[int]], default: None ) –

    Channels for the DynamicTNet local (point-wise) MLP.

  • stnet_global_channels (Optional[Sequence[int]], default: None ) –

    Channels for the DynamicTNet global MLP.

  • stnet_num_neighbors (int, default: 20 ) –

    Number of kNN neighbors for the DynamicTNet.

  • proj_channels (int, default: 1024 ) –

    Number of channels for the projection layer after the encoder.

  • channels (Sequence[Union[int, Sequence[int]]]) –

    List of channels for each encoder block.

  • head_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    List of channels for each head block.

  • num_neighbors (Union[int, Sequence[int]]) –

    Maximum number of neighbors for each encoder block.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the activation function.

  • act_first (bool, default: False ) –

    Whether to apply activation before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the normalization layer.

  • bias (bool, default: True ) –

    Whether to use bias in linear / MLP layers.

  • dropout (float, default: 0.0 ) –

    Dropout probability in the decoder head.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: encoder features, global feature and category embedding.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: encoder features, global feature and category embedding.

forward_decoder [source] ΒΆ

forward_decoder(*args: Any, **kwargs: Any) -> Any

Decode encoder features back to per-point resolution.

Canonical signature: forward_decoder(x, ..., intermediates), consuming the output of forward_features and returning per-point features \((N, C)\). Models whose encoder already emits per-point features (DGCNN, PointNet) have no decoder and raise NotImplementedError.

configure_stnet [source] ΒΆ

configure_stnet() -> Optional[DynamicTNet]

Build the dynamic spatial transformer, or None when stnet_edge_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> DGCNNEncoder

Build the DGCNNEncoder backbone.

configure_proj [source] ΒΆ

configure_proj() -> MLP

Build the projection MLP producing the global feature.

configure_cat_embed [source] ΒΆ

configure_cat_embed() -> MLP

Build the MLP embedding the one-hot shape category.

DGCNNSegmentation [source] ΒΆ

DGCNNSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stnet_local_channels: Optional[Sequence[int]] = None,
    stnet_global_channels: Optional[Sequence[int]] = None,
    proj_channels: int = 1024,
    channels: Sequence[Union[int, Sequence[int]]],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    num_neighbors: Union[int, Sequence[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,
)

Bases: SemanticSegmentationModel

Semantic segmentation model as described in the paper "Dynamic Graph CNN for Learning on Point Clouds" by Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, Justin M. Solomon.

DGCNN introduces the EdgeConv operator, which computes features on dynamically constructed k-nearest neighbor (k-NN) graphs at each layer. Graphs are recomputed in the learned feature space, allowing the network to capture both local geometric relationships and long-range semantic structures.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

  • stnet_local_channels (Optional[Sequence[int]], default: None ) –

    List of channels for the local spatial transformer network. If None, the spatial transformer network is not used.

  • stnet_global_channels (Optional[Sequence[int]], default: None ) –

    List of channels for the global spatial transformer network. If None, the spatial transformer network is not used.

  • proj_channels (int, default: 1024 ) –

    Number of channels for the projection layer after the encoder.

  • channels (Sequence[Union[int, Sequence[int]]]) –

    List of channels for each encoder block.

  • head_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    List of channels for each head block.

  • num_neighbors (Union[int, Sequence[int]]) –

    Maximum number of neighbors for each encoder block.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the activation function.

  • act_first (bool, default: False ) –

    Whether to apply activation before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the normalization layer.

  • bias (bool, default: True ) –

    Whether to use bias in linear / MLP layers.

  • dropout (float, default: 0.0 ) –

    Dropout probability before the classification head.

Methods:

  • forward_decoder –

    Decode encoder features back to per-point resolution.

  • configure_stnet –

    Build the spatial transformer aligning the input coordinates, or None when its channels are unset.

  • configure_encoder –

    Build the DGCNNEncoder backbone.

  • configure_proj –

    Build the projection MLP producing the global feature.

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: encoder features plus the broadcast global feature.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: encoder features plus the broadcast global feature.

forward_decoder [source] ΒΆ

forward_decoder(*args: Any, **kwargs: Any) -> Any

Decode encoder features back to per-point resolution.

Canonical signature: forward_decoder(x, ..., intermediates), consuming the output of forward_features and returning per-point features \((N, C)\). Models whose encoder already emits per-point features (DGCNN, PointNet) have no decoder and raise NotImplementedError.

configure_stnet [source] ΒΆ

configure_stnet() -> Optional[TNet]

Build the spatial transformer aligning the input coordinates, or None when its channels are unset.

configure_encoder [source] ΒΆ

configure_encoder() -> DGCNNEncoder

Build the DGCNNEncoder backbone.

configure_proj [source] ΒΆ

configure_proj() -> MLP

Build the projection MLP producing the global feature.

KPFCNNClassification [source] ΒΆ

KPFCNNClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stem_channels: Optional[int] = None,
    stem_type: Literal["linear", "kpconv"] = "kpconv",
    encoder_depths: Sequence[int],
    encoder_channels: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    grid_sizes: Sequence[float],
    radii: Sequence[float],
    kernel_size: int,
    kp_radius: Union[float, Sequence[float]],
    kp_sigma: Union[float, Sequence[float]],
    kp_influence: str = "linear",
    fixed_position: Literal[
        "none", "center", "vertical"
    ] = "center",
    aggregation_mode: str = "sum",
    deformable: Union[bool, Sequence] = False,
    modulated: Union[bool, Sequence] = 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,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
)

Bases: ClassificationModel

KPConv Network for classification tasks as described in the paper KPConv: Flexible and Efficient Convolution for Point Clouds by Hugues Thomas, Charles R. Qi, Jean-Emmanuel Deschaud, Beatriz Marcotegui, FranΓ§ois Goulette, Leonidas J. Guibas.

KPConv introduces a novel point convolution operator that uses kernel points to define the spatial extent and weights of the convolution. The kernel points are arranged in space to define the convolution pattern, with weights determined by their spatial correlation with input points. This allows for flexible and efficient convolution on irregular point clouds while maintaining permutation invariance and translation invariance. The network uses a hierarchical architecture with strided convolutions for spatial pooling and feature aggregation.

Note

The implementation is based on the original paper and the authors' code KPConv-PyTorch.

Important

This implementation was completely rewritten to be compatible with torch-geometric library.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

  • stem_channels (Optional[int], default: None ) –

    Number of channels in the stem layer.

  • stem_type (Literal['linear', 'kpconv'], default: 'kpconv' ) –

    Type of stem layer to use.

  • encoder_depths (Sequence[int]) –

    List of depths for each encoder block, i.e. corresponds to the number of residual blocks at each level.

  • encoder_channels (Sequence[int]) –

    List of channels for each encoder block.

  • encoder_num_neighbors (Sequence[int]) –

    List of maximum number of neighbors for each encoder block.

  • grid_sizes (Sequence[float]) –

    List of grid sizes for each downsampling block.

  • radii (Sequence[float]) –

    Search radius for each downsampling block.

  • kernel_size (int) –

    Size of the kernel for each KPConv block.

  • kp_radius (Union[float, Sequence[float]]) –

    List of kernel radius for KPConv blocks, at each level.

  • kp_sigma (Union[float, Sequence[float]]) –

    List of kernel extent for KPConv blocks, at each level.

  • kp_influence (str, default: 'linear' ) –

    Influence function to use for KPConv blocks. Options are "constant", "linear", "gaussian".

  • fixed_position (Literal['none', 'center', 'vertical'], default: 'center' ) –

    Whether to fix the position of the kernel points in KPConv blocks. Options are "none", "center", "vertical".

  • aggregation_mode (str, default: 'sum' ) –

    Aggregation mode to use for the KPConv blocks. Options are "sum", "mean", "max".

  • deformable (Union[bool, Sequence], default: False ) –

    Whether to use a deformable kernel for the KPConv blocks.

  • modulated (Union[bool, Sequence], default: False ) –

    Whether to use a modulated kernel in KPConv operation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization to use for the KPConv blocks.

  • act (Union[str, Callable, None], default: 'leaky_relu' ) –

    Activation function to use for the KPConv blocks.

  • bias (bool, default: False ) –

    Whether to use a bias for the KPConv blocks.

  • dropout (float, default: 0.0 ) –

    Dropout rate before the classification head.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling method to use before the classification head. Options are "max", "mean".

Methods:

  • configure_stem –

    Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

  • configure_encoder –

    Build the KPFCNNEncoder backbone.

  • reset_classifier –

    Resets the classification head with new parameters.

  • forward_features –

    Forward pass of the encoder, returning pre-pooling features.

  • forward_head –

    Forward pass of the classification head from pre-pooling features.

  • forward –

    Forward pass of the classification model.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_stem [source] ΒΆ

configure_stem() -> Optional[Module]

Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> KPFCNNEncoder

Build the KPFCNNEncoder backbone.

reset_classifier [source] ΒΆ

reset_classifier(
    num_classes: int,
    global_pool: Optional[PoolLike] = None,
    **kwargs: Any,
) -> None

Resets the classification head with new parameters.

Note

To set an empty classification head, use num_classes=0.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • global_pool (Optional[PoolLike], default: None ) –

    Pooling method to aggregate point features ("max" or "mean"). If None, keeps the current pooling.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the classification head.

forward_features [source] ΒΆ

forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: Literal[True],
) -> Tuple[
    Tensor, Tensor, Tensor, List[PooledFeaturesDict]
]
forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: Literal[False] = False,
) -> Tuple[Tensor, Tensor, Tensor]
forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: bool = False,
) -> Any

Forward pass of the encoder, returning pre-pooling features.

Parameters:

  • x (OptTensor) –

    Point features of shape \((N, C)\). If None, pos is used as features, which requires in_channels to match the dimension of pos.

  • pos (Tensor) –

    Point coordinates of shape \((N, D)\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • return_intermediates (bool, default: False ) –

    Whether to also return the per-stage intermediates.

Returns:

  • Any –

    A tuple (x, pos, batch) at the coarsest level, where x has shape \((N', \text{encoder\_channels}[-1])\) and \(N'\) is the number of downsampled points. If return_intermediates=True, the per-stage intermediates are appended to the tuple.

forward_head [source] ΒΆ

forward_head(
    x: Tensor, batch: Tensor, pre_logits: bool = False
) -> Tensor

Forward pass of the classification head from pre-pooling features.

Parameters:

  • x (Tensor) –

    Pre-pooling features of shape \((N', \text{encoder\_channels}[-1])\) where \(N'\) is the number of downsampled points.

  • batch (Tensor) –

    Batch indices for each downsampled point of shape \((N',)\).

  • pre_logits (bool, default: False ) –

    Whether to return pre-logits. Defaults to False.

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

forward [source] ΒΆ

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

Forward pass of the classification model.

Parameters:

  • x (OptTensor) –

    Point features of shape \((N, C)\). If None, pos is used as features, which requires in_channels to match the dimension of pos.

  • pos (Tensor) –

    Point coordinates of shape \((N, D)\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

KPFCNNSegmentation [source] ΒΆ

KPFCNNSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stem_channels: Optional[int] = None,
    stem_type: Literal["linear", "kpconv"] = "kpconv",
    encoder_depths: Sequence[int],
    encoder_channels: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    fp_channels: Sequence[Sequence[int]],
    grid_sizes: Sequence[float],
    radii: Sequence[float],
    kernel_size: int,
    kp_radius: Union[float, Sequence[float]],
    kp_sigma: Union[float, Sequence[float]],
    kp_influence: str = "linear",
    fixed_position: Literal[
        "none", "center", "vertical"
    ] = "center",
    aggregation_mode: str = "sum",
    deformable: Union[bool, Sequence] = False,
    modulated: Union[bool, Sequence] = 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,
    dropout: float = 0.0,
    head_channels: Optional[Sequence[int]] = None,
)

Bases: SemanticSegmentationModel

KPConv Network for segmentation tasks as described in the paper KPConv: Flexible and Efficient Convolution for Point Clouds by Hugues Thomas, Charles R. Qi, Jean-Emmanuel Deschaud, Beatriz Marcotegui, FranΓ§ois Goulette, Leonidas J. Guibas.

KPConv introduces a novel point convolution operator that uses kernel points to define the spatial extent and weights of the convolution. The kernel points are arranged in space to define the convolution pattern, with weights determined by their spatial correlation with input points. This allows for flexible and efficient convolution on irregular point clouds while maintaining permutation invariance and translation invariance. The network uses a hierarchical architecture with strided convolutions for spatial pooling and feature aggregation.

Note

The implementation is based on the original paper and the authors' code KPConv-PyTorch.

Important

This implementation was completely rewritten to be compatible with torch-geometric library.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

  • stem_channels (Optional[int], default: None ) –

    Number of channels in the stem layer.

  • stem_type (Literal['linear', 'kpconv'], default: 'kpconv' ) –

    Type of stem layer to use.

  • encoder_depths (Sequence[int]) –

    List of depths for each encoder block, i.e. corresponds to the number of residual blocks at each level.

  • encoder_channels (Sequence[int]) –

    List of channels for each encoder block.

  • encoder_num_neighbors (Sequence[int]) –

    List of maximum number of neighbors for each encoder block.

  • fp_channels (Sequence[Sequence[int]]) –

    List of channels for each feature propagation block.

  • grid_sizes (Sequence[float]) –

    List of grid sizes for each downsampling block.

  • radii (Sequence[float]) –

    Search radius for each downsampling block.

  • kernel_size (int) –

    Size of the kernel for each KPConv block.

  • kp_radius (Union[float, Sequence[float]]) –

    List of kernel radius for KPConv blocks, at each level.

  • kp_sigma (Union[float, Sequence[float]]) –

    List of kernel extent for KPConv blocks, at each level.

  • kp_influence (str, default: 'linear' ) –

    Influence function to use for KPConv blocks. Options are "constant", "linear", "gaussian".

  • fixed_position (Literal['none', 'center', 'vertical'], default: 'center' ) –

    Whether to fix the position of the kernel points in KPConv blocks. Options are "none", "center", "vertical".

  • aggregation_mode (str, default: 'sum' ) –

    Aggregation mode to use for the KPConv blocks. Options are "sum", "mean", "max".

  • deformable (Union[bool, Sequence], default: False ) –

    Whether to use a deformable kernel for the KPConv blocks.

  • modulated (Union[bool, Sequence], default: False ) –

    Whether to use a modulated kernel in KPConv operation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization to use for the KPConv blocks.

  • act (Union[str, Callable, None], default: 'leaky_relu' ) –

    Activation function to use for the KPConv blocks.

  • bias (bool, default: False ) –

    Whether to use a bias for the KPConv blocks.

  • dropout (float, default: 0.0 ) –

    Dropout rate before the classification head.

Methods:

  • configure_stem –

    Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

  • configure_encoder –

    Build the KPFCNNEncoder backbone.

  • configure_decoder –

    Build the PointNet2Decoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_stem [source] ΒΆ

configure_stem() -> Optional[Module]

Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> KPFCNNEncoder

Build the KPFCNNEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointNet2Decoder

Build the PointNet2Decoder upsampling the coarsest features back through the encoder skips.

LIONDetection [source] ΒΆ

LIONDetection(
    in_channels: int = 5,
    num_classes: int = 10,
    *,
    voxel_size: Sequence[float] = (0.3, 0.3, 0.25),
    point_cloud_range: Sequence[float] = (
        -54.0,
        -54.0,
        -5.0,
        54.0,
        54.0,
        3.0,
    ),
    channels: int = 128,
    vfe_num_filters: Sequence[int] = (128, 128),
    depths: Sequence[int] = (2, 2, 2, 2),
    window_shape: Sequence[Sequence[int]] = (
        (13, 13, 32),
        (13, 13, 16),
        (13, 13, 8),
        (13, 13, 4),
    ),
    group_size: Sequence[int] = (4096, 2048, 1024, 512),
    diffusion: bool = True,
    diff_scale: float = 0.2,
    layer_nums: Sequence[int] = (1, 2, 2),
    layer_strides: Sequence[int] = (1, 2, 2),
    num_filters: Sequence[int] = (128, 128, 256),
    upsample_strides: Sequence[float] = (0.5, 1, 2),
    num_upsample_filters: Sequence[int] = (128, 128, 128),
    feature_map_stride: int = 2,
    local_max_classes: Sequence[int] = (),
    d_state: int = 16,
    d_conv: int = 4,
    expand: int = 2,
)

Bases: DetectionModel

LION: linear group RNN (Mamba) 3D object detector (packed point format).

Reference: LION: Linear Group RNN for 3D Object Detection in Point Clouds (Liu et al., 2024). Reference implementation: happinesslz/LION.

Points are encoded into voxels by a dynamic mean VFE, processed by a hierarchical sparse backbone that serializes voxels into spatially grouped windows and runs a bidirectional Mamba operator (the linear group RNN), with periodic 3D voxel-generation / height-merging steps. The one-stride output is folded to a dense BEV map, refined by a residual 2D backbone, and decoded by a query-based TransFusionHead.

Parameters:

  • in_channels (int, default: 5 ) –

    Raw point feature channels including xyz (5 for nuScenes \(x, y, z, \text{intensity}, \Delta t\)).

  • num_classes (int, default: 10 ) –

    Number of foreground classes (10 for nuScenes).

  • voxel_size (Sequence[float], default: (0.3, 0.3, 0.25) ) –

    Voxel size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (-54.0, -54.0, -5.0, 54.0, 54.0, 3.0) ) –

    Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).

  • channels (int, default: 128 ) –

    Backbone / VFE feature channels.

  • vfe_num_filters (Sequence[int], default: (128, 128) ) –

    PFN widths of the dynamic mean VFE.

  • depths (Sequence[int], default: (2, 2, 2, 2) ) –

    Per-stage encoder/decoder depth of the 3D backbone.

  • window_shape (Sequence[Sequence[int]], default: ((13, 13, 32), (13, 13, 16), (13, 13, 8), (13, 13, 4)) ) –

    Per-stage window size \((w_x, w_y, w_z)\).

  • group_size (Sequence[int], default: (4096, 2048, 1024, 512) ) –

    Per-stage Mamba group length.

  • diffusion (bool, default: True ) –

    Enable voxel-generation in the height-merge steps.

  • diff_scale (float, default: 0.2 ) –

    Fraction of voxels expanded by diffusion.

  • layer_nums (Sequence[int], default: (1, 2, 2) ) –

    2D backbone residual-block counts per level.

  • layer_strides (Sequence[int], default: (1, 2, 2) ) –

    2D backbone downsample strides per level.

  • num_filters (Sequence[int], default: (128, 128, 256) ) –

    2D backbone channel widths per level.

  • upsample_strides (Sequence[float], default: (0.5, 1, 2) ) –

    2D backbone upsample factors per level.

  • num_upsample_filters (Sequence[int], default: (128, 128, 128) ) –

    2D backbone upsample channels per level.

  • feature_map_stride (int, default: 2 ) –

    BEV stride of the head.

  • local_max_classes (Sequence[int], default: () ) –

    Crowded small-object class indices passed to the head (nuScenes pedestrian / traffic-cone); see TransFusionHead.

  • d_state (int, default: 16 ) –

    SSM state width.

  • d_conv (int, default: 4 ) –

    Causal-conv kernel width.

  • expand (int, default: 2 ) –

    Inner-width expansion factor of the Mamba operator.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_vfe –

    Build the dynamic mean voxel feature encoder.

  • configure_backbone_3d –

    Build the sparse 3D linear group RNN backbone.

  • configure_backbone –

    Build the residual 2D BEV backbone.

  • configure_head –

    Build the query-based TransFusion detection head.

  • decode –

    Decode a forward output into raw candidate detections (see TransFusionHead.decode).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the BEV feature map entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_vfe [source] ΒΆ

configure_vfe() -> DynamicMeanVFE

Build the dynamic mean voxel feature encoder.

configure_backbone_3d [source] ΒΆ

configure_backbone_3d() -> LION3DBackbone

Build the sparse 3D linear group RNN backbone.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVResidualBackbone

Build the residual 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> TransFusionHead

Build the query-based TransFusion detection head.

decode [source] ΒΆ

Decode a forward output into raw candidate detections (see TransFusionHead.decode).

OctFormerClassification [source] ΒΆ

OctFormerClassification(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Union[int, Sequence[int]],
    encoder_channels: Sequence[int],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    encoder_depths: Sequence[int],
    num_heads: Sequence[int],
    patch_size: int = 26,
    dilation: int = 4,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.5,
    nempty: bool = True,
    use_checkpoint: bool = True,
    use_rpe: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    head_act: Union[str, Callable, None] = "gelu",
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    dropout: float = 0.0,
    global_pool: PoolLike = "mean",
)

Bases: ClassificationModel

OctFormer classification model from OctFormer: Octree-based Transformers for 3D Point Clouds by Peng-Shuai Wang.

An octree convolution stem embeds the input signal, then attention runs over patches of equal point count sorted along the octree's space-filling curve, one stage per octree depth. Point features are pooled globally after the encoder for classification.

Methods:

  • configure_stem –

    Build the OctreePatchEmbed stem.

  • configure_encoder –

    Build the OctFormerEncoder backbone.

  • get_encoder_depth –

    Octree depth at which the first encoder stage runs, once the stem downsamplings are accounted for.

  • get_head_depth –

    Octree depth of the encoder output, which is the depth the head pools over.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_stem [source] ΒΆ

configure_stem() -> Module

Build the OctreePatchEmbed stem.

configure_encoder [source] ΒΆ

configure_encoder() -> Module

Build the OctFormerEncoder backbone.

get_encoder_depth [source] ΒΆ

get_encoder_depth(depth: int) -> int

Octree depth at which the first encoder stage runs, once the stem downsamplings are accounted for.

Parameters:

  • depth (int) –

    Octree depth of the input signal.

Returns:

  • int –

    The octree depth of the encoder's first stage.

get_head_depth [source] ΒΆ

get_head_depth(depth: int) -> int

Octree depth of the encoder output, which is the depth the head pools over.

Parameters:

  • depth (int) –

    Octree depth of the input signal.

Returns:

  • int –

    The octree depth of the encoder's last stage.

OctFormerSegmentation [source] ΒΆ

OctFormerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Union[int, Sequence[int]],
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    num_heads: Sequence[int],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    fpn_channels: int,
    patch_size: int = 26,
    dilation: int = 4,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.5,
    nempty: bool = True,
    use_checkpoint: bool = True,
    use_rpe: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    head_act: Union[str, Callable, None] = "gelu",
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    dropout: float = 0.5,
)

Bases: SemanticSegmentationModel

OctFormer segmentation model from OctFormer: Octree-based Transformers for 3D Point Clouds by Peng-Shuai Wang.

The octree attention encoder is followed by a feature pyramid decoder that merges every stage at fpn_channels and upsamples back to the input octree depth, then a per-point MLP head.

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the features entering the head.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the features entering the head.

configure_stem [source] ΒΆ

configure_stem() -> Module

Build the OctreePatchEmbed stem.

configure_encoder [source] ΒΆ

configure_encoder() -> Module

Build the OctFormerEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> Module

Build the OctFormerDecoder, with one upsampling per stem downsampling.

PointBERTClassification [source] ΒΆ

PointBERTClassification(
    in_channels: int,
    num_classes: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_dims: int = 256,
    token_local_channels: Sequence[int] = (128, 256),
    token_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    drop_path: float = 0.1,
    spatial_dim: int = 3,
    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,
    head_act: Union[str, Callable, None] = "relu",
    dropout: float = 0.5,
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = 256,
)

Bases: ClassificationModel

Point-BERT classification model.

Implements the finetuning model of Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling, adapted from lulutang0608/Point-BERT.

A PointBERTEncoder backbone followed by a 2-layer MLP head. The global feature concatenates the class token with the max-pooled patch tokens, so the head input dimension is \(2d\).

Parameters:

  • in_channels (int) –

    The number of per-point feature channels concatenated to the coordinates (\(0\) for coordinates only).

  • num_classes (int) –

    The number of output classes.

  • embed_dim (int, default: 384 ) –

    The transformer dimension \(d\).

  • depth (int, default: 12 ) –

    The number of transformer blocks.

  • num_heads (int, default: 6 ) –

    The number of attention heads.

  • num_groups (int, default: 64 ) –

    The number of patches \(G\).

  • group_size (int, default: 32 ) –

    The number of points \(M\) per patch.

  • encoder_dims (int, default: 256 ) –

    The token-embedding dimension before the linear bridge.

  • token_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the tokenizer's per-point MLP.

  • token_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the tokenizer's per-patch MLP.

  • pos_embed_channels (Sequence[int], default: (128,) ) –

    Hidden widths of the positional-embedding MLP.

  • drop_path (float, default: 0.1 ) –

    The stochastic depth rate (identity at eval).

  • spatial_dim (int, default: 3 ) –

    The number of spatial dimensions of the coordinates.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    The activation used in the transformer MLPs.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: LayerNorm ) –

    The normalization used in the transformer blocks.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization.

  • head_act (Union[str, Callable, None], default: 'relu' ) –

    The activation used in the classification head.

  • dropout (float, default: 0.5 ) –

    The dropout rate of the classification head.

  • head_channels (Optional[Union[int, Sequence[int]]], default: 256 ) –

    The hidden width(s) of the classification head.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: \((B, \text{num\_classes})\).
Example
import torch
from torch_pointcloud.models.point_bert import PointBERTClassification

model = PointBERTClassification(in_channels=0, num_classes=40)
pos = torch.randn(2048, 3)
batch = torch.cat([torch.zeros(1024), torch.ones(1024)]).long()
logits = model(None, pos, batch)
print(logits.shape)

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointBERTEncoder

Build the PointBERTEncoder backbone.

PointBERTDiscreteVAE [source] ΒΆ

PointBERTDiscreteVAE(
    in_channels: int,
    *,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_dims: int = 256,
    token_local_channels: Sequence[int] = (128, 256),
    token_global_channels: Sequence[int] = (512,),
    num_tokens: int = 8192,
    tokens_dims: int = 256,
    decoder_dims: int = 256,
    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: PretrainingModel

Point-BERT discrete VAE point tokenizer.

Implements the dVAE of Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling, adapted from lulutang0608/Point-BERT.

The dVAE produces the discrete point tokens used as the masked-modeling targets in Point-BERT pretraining: a mini-PointNet token encoder, a DGCNN that maps tokens to codebook logits, a learned codebook, a second DGCNN, and a folding decoder reconstructing the patches.

Parameters:

  • in_channels (int) –

    The number of input channels (\(0\), coordinates only).

  • num_groups (int, default: 64 ) –

    The number of patches \(G\).

  • group_size (int, default: 32 ) –

    The number of points \(M\) per patch.

  • encoder_dims (int, default: 256 ) –

    The mini-PointNet token-embedding dimension.

  • token_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the tokenizer's per-point MLP.

  • token_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the tokenizer's per-patch MLP.

  • num_tokens (int, default: 8192 ) –

    The codebook vocabulary size.

  • tokens_dims (int, default: 256 ) –

    The codebook embedding dimension.

  • decoder_dims (int, default: 256 ) –

    The dimension fed to the folding decoder.

  • act (Union[str, Callable, None], default: 'relu' ) –

    The activation used in the token encoder and the folding decoder.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the token / decoder activation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    The normalization used in the token encoder and the folding decoder.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the token / decoder normalization.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: a dict with logits \((B, G, \text{num\_tokens})\) and reconstructions.
Example
import torch
from torch_pointcloud.models.point_bert import PointBERTDiscreteVAE

model = PointBERTDiscreteVAE(in_channels=0)
pos = torch.randn(2048, 3)
batch = torch.cat([torch.zeros(1024), torch.ones(1024)]).long()
out = model(None, pos, batch)
print(out["logits"].shape, out["fine"].shape)

Methods:

  • configure_encoder –

    Build the mini-PointNet token encoder embedding each patch.

  • configure_dgcnn_1 –

    Build the DGCNN mapping patch tokens to codebook logits.

  • configure_dgcnn_2 –

    Build the DGCNN mapping sampled codebook embeddings to decoder features.

  • configure_decoder –

    Build the folding decoder reconstructing each patch.

  • tokenize –

    Return the codebook logits of every group, without sampling or decoding them.

configure_encoder [source] ΒΆ

configure_encoder() -> PointPatchEmbed

Build the mini-PointNet token encoder embedding each patch.

configure_dgcnn_1 [source] ΒΆ

configure_dgcnn_1() -> TokenDGCNN

Build the DGCNN mapping patch tokens to codebook logits.

configure_dgcnn_2 [source] ΒΆ

configure_dgcnn_2() -> TokenDGCNN

Build the DGCNN mapping sampled codebook embeddings to decoder features.

configure_decoder [source] ΒΆ

configure_decoder() -> FoldingDecoder

Build the folding decoder reconstructing each patch.

tokenize [source] ΒΆ

tokenize(pos: Tensor, batch: Tensor) -> Tensor

Return the codebook logits of every group, without sampling or decoding them.

Parameters:

  • pos (Tensor) –

    Point coordinates of shape \((N, 3)\).

  • batch (Tensor) –

    Per-point batch index of shape \((N,)\).

Returns:

  • Tensor –

    Logits over the codebook of shape \((B, G, \text{num\_tokens})\).

PointBERTPretraining [source] ΒΆ

PointBERTPretraining(
    in_channels: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_dims: int = 256,
    token_local_channels: Sequence[int] = (128, 256),
    token_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    num_tokens: int = 8192,
    cls_dim: int = 512,
    mask_ratio: Tuple[float, float] = (0.25, 0.45),
    drop_path: float = 0.1,
    spatial_dim: int = 3,
    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,
    token_act: Union[str, Callable, None] = "relu",
    token_act_kwargs: Optional[Dict[str, Any]] = None,
    token_norm: Union[str, Callable, None] = "batch_norm",
    token_norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: PretrainingModel

Point-BERT masked point modeling backbone (pretrain).

Implements the masked transformer (transformer_q) of Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling, adapted from lulutang0608/Point-BERT.

It embeds patches and, in training mode, replaces a contiguous block of patch tokens (a ratio drawn uniformly from mask_ratio, centered on a random patch) with a learned mask token while keeping their positional embeddings. It then runs the transformer and exposes a token-classification head (lm_head, predicting dVAE codebook ids) and a contrastive class head (cls_head). In eval mode no tokens are masked. The MoCo / cutmix machinery of the full pretraining objective is omitted; this module is the reusable encoder that downstream finetuning loads.

Parameters:

  • in_channels (int) –

    The number of per-point feature channels concatenated to the coordinates (\(0\) for coordinates only).

  • embed_dim (int, default: 384 ) –

    The transformer dimension \(d\).

  • depth (int, default: 12 ) –

    The number of transformer blocks.

  • num_heads (int, default: 6 ) –

    The number of attention heads.

  • num_groups (int, default: 64 ) –

    The number of patches \(G\).

  • group_size (int, default: 32 ) –

    The number of points \(M\) per patch.

  • encoder_dims (int, default: 256 ) –

    The token-embedding dimension before the linear bridge.

  • token_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the tokenizer's per-point MLP.

  • token_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the tokenizer's per-patch MLP.

  • pos_embed_channels (Sequence[int], default: (128,) ) –

    Hidden widths of the positional-embedding MLP.

  • num_tokens (int, default: 8192 ) –

    The dVAE vocabulary size predicted by lm_head.

  • cls_dim (int, default: 512 ) –

    The contrastive head output dimension.

  • mask_ratio (Tuple[float, float], default: (0.25, 0.45) ) –

    Lower / upper bounds of the block-masking ratio.

  • drop_path (float, default: 0.1 ) –

    The stochastic depth rate (identity at eval).

  • spatial_dim (int, default: 3 ) –

    The number of spatial dimensions of the coordinates.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    The activation used in the transformer MLPs and the contrastive head.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: LayerNorm ) –

    The normalization used in the transformer blocks.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization.

  • token_act (Union[str, Callable, None], default: 'relu' ) –

    The activation used in the token encoder.

  • token_act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the token-encoder activation.

  • token_norm (Union[str, Callable, None], default: 'batch_norm' ) –

    The normalization used in the token encoder.

  • token_norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the token-encoder normalization.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: a dict with cls_feature \((B, \text{cls\_dim})\) and logits \((B, G, \text{num\_tokens})\).
Example
import torch
from torch_pointcloud.models.point_bert import PointBERTPretraining

model = PointBERTPretraining(in_channels=0)
pos = torch.randn(2048, 3)
batch = torch.cat([torch.zeros(1024), torch.ones(1024)]).long()
out = model(None, pos, batch)
print(out["cls_feature"].shape, out["logits"].shape)

Methods:

  • configure_encoder –

    Build the mini-PointNet token encoder embedding each patch.

  • configure_reduce_dim –

    Build the linear bridge from the token-embedding dimension to the transformer dimension.

  • configure_pos_embed –

    Build the positional-embedding MLP mapping patch centers to the transformer dimension.

  • configure_blocks –

    Build the transformer blocks with a linearly scaled stochastic-depth schedule.

  • configure_lm_head –

    Build the token-classification head predicting dVAE codebook ids.

  • configure_cls_head –

    Build the contrastive class-token head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointPatchEmbed

Build the mini-PointNet token encoder embedding each patch.

configure_reduce_dim [source] ΒΆ

configure_reduce_dim() -> Linear

Build the linear bridge from the token-embedding dimension to the transformer dimension.

configure_pos_embed [source] ΒΆ

configure_pos_embed() -> MLP

Build the positional-embedding MLP mapping patch centers to the transformer dimension.

configure_blocks [source] ΒΆ

configure_blocks() -> ModuleList

Build the transformer blocks with a linearly scaled stochastic-depth schedule.

configure_lm_head [source] ΒΆ

configure_lm_head() -> Linear

Build the token-classification head predicting dVAE codebook ids.

configure_cls_head [source] ΒΆ

configure_cls_head() -> MLP

Build the contrastive class-token head.

PointM2AEClassification [source] ΒΆ

PointM2AEClassification(
    in_channels: int,
    num_classes: int,
    *,
    group_sizes: Sequence[int] = (16, 8, 8),
    num_groups: Sequence[int] = (512, 256, 64),
    encoder_depths: Sequence[int] = (5, 5, 5),
    encoder_dims: Sequence[int] = (96, 192, 384),
    token_local_channels: Sequence[int] = (128, 256),
    token_global_channels: Sequence[int] = (512,),
    local_radius: Sequence[float] = (0.32, 0.64, 1.28),
    num_heads: int = 6,
    drop_path: float = 0.1,
    concat_pooling: bool = False,
    dropout: float = 0.5,
    head_channels: Sequence[int] = (256, 256),
)

Bases: ClassificationModel

Implementation of the Point-M2AE classification model.

Point-M2AE: Multi-scale Masked Autoencoders for Hierarchical Point Cloud Pre-training. This implementation is adapted from the official repository ZrrSkywalker/Point-M2AE.

The model groups the input cloud at multiple scales, encodes it with the hierarchical transformer encoder and pools the finest-stage tokens into a feature passed to an MLP head. The pooling matches the reference: ModelNet40 sums the token mean and max, while ScanObjectNN concatenates the mean over all tokens with the max over the tokens after the first one.

Parameters:

  • in_channels (int) –

    The number of per-point feature channels concatenated to the coordinates (\(0\) for coordinates only).

  • num_classes (int) –

    The number of output classes.

  • group_sizes (Sequence[int], default: (16, 8, 8) ) –

    The neighborhood size per stage.

  • num_groups (Sequence[int], default: (512, 256, 64) ) –

    The number of centers per stage.

  • encoder_depths (Sequence[int], default: (5, 5, 5) ) –

    The number of transformer blocks per stage.

  • encoder_dims (Sequence[int], default: (96, 192, 384) ) –

    The channel width per stage.

  • token_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the first-stage token embedder's per-point MLP.

  • token_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the first-stage token embedder's per-group MLP.

  • local_radius (Sequence[float], default: (0.32, 0.64, 1.28) ) –

    The local-attention radius per stage.

  • num_heads (int, default: 6 ) –

    The number of attention heads.

  • drop_path (float, default: 0.1 ) –

    The maximum stochastic depth rate.

  • concat_pooling (bool, default: False ) –

    Use the ScanObjectNN concat pooling when True, the ModelNet40 sum pooling otherwise.

  • dropout (float, default: 0.5 ) –

    The dropout rate in the classification head.

  • head_channels (Sequence[int], default: (256, 256) ) –

    The hidden widths of the classification head.

Shape
  • pos: \((N, 3)\)
  • batch: \((N,)\)
  • Output: \((B, C)\) where \(C\) is the number of classes.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_h_encoder [source] ΒΆ

configure_h_encoder() -> HierarchicalEncoder

Build the hierarchical transformer encoder.

PointM2AEPartSegmentation [source] ΒΆ

PointM2AEPartSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    num_categories: int = 16,
    group_sizes: Sequence[int] = (16, 8, 8),
    num_groups: Sequence[int] = (512, 256, 64),
    encoder_depths: Sequence[int] = (5, 5, 5),
    encoder_dims: Sequence[int] = (96, 192, 384),
    token_local_channels: Sequence[int] = (128, 256),
    token_global_channels: Sequence[int] = (512,),
    local_radius: Sequence[float] = (0.32, 0.64, 1.28),
    num_heads: int = 6,
)

Bases: PartSegmentationModel

Implementation of the Point-M2AE part-segmentation model.

Point-M2AE: Multi-scale Masked Autoencoders for Hierarchical Point Cloud Pre-training. This implementation is adapted from the official repository ZrrSkywalker/Point-M2AE.

Per-stage encoder features are propagated back to the full-resolution cloud, concatenated with a global feature and the one-hot object label, and decoded into per-point logits. Every sample in the packed batch must contain the same number of points; a ragged batch raises ValueError.

Parameters:

  • in_channels (int) –

    The number of per-point feature channels concatenated to the coordinates (\(0\) for coordinates only).

  • num_classes (int) –

    The number of part-segmentation classes.

  • num_categories (int, default: 16 ) –

    The number of object categories (for the one-hot label embedding).

  • group_sizes (Sequence[int], default: (16, 8, 8) ) –

    The neighborhood size per stage.

  • num_groups (Sequence[int], default: (512, 256, 64) ) –

    The number of centers per stage.

  • encoder_depths (Sequence[int], default: (5, 5, 5) ) –

    The number of transformer blocks per stage.

  • encoder_dims (Sequence[int], default: (96, 192, 384) ) –

    The channel width per stage.

  • token_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the first-stage token embedder's per-point MLP.

  • token_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the first-stage token embedder's per-group MLP.

  • local_radius (Sequence[float], default: (0.32, 0.64, 1.28) ) –

    The local-attention radius per stage.

  • num_heads (int, default: 6 ) –

    The number of attention heads.

Shape
  • pos: \((N, 3)\)
  • batch: \((N,)\)
  • category: \((B, \text{num\_categories})\) one-hot object label
  • Output: \((N, C)\) where \(C\) is the number of part classes.

Methods:

  • configure_h_encoder –

    Build the hierarchical transformer encoder with per-stage output norms.

  • configure_label_conv –

    Build the MLP embedding the category one-hot for the global branch.

  • configure_propagations –

    Build the per-stage feature-propagation modules interpolating stage features to every point.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the features entering the head.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the features entering the head.

configure_h_encoder [source] ΒΆ

configure_h_encoder() -> HierarchicalEncoder

Build the hierarchical transformer encoder with per-stage output norms.

configure_label_conv [source] ΒΆ

configure_label_conv() -> MLP

Build the MLP embedding the category one-hot for the global branch.

configure_propagations [source] ΒΆ

configure_propagations() -> ModuleList

Build the per-stage feature-propagation modules interpolating stage features to every point.

PointM2AEPretraining [source] ΒΆ

PointM2AEPretraining(
    in_channels: int,
    *,
    group_sizes: Sequence[int] = (16, 8, 8),
    num_groups: Sequence[int] = (512, 256, 64),
    mask_ratio: float = 0.8,
    encoder_depths: Sequence[int] = (5, 5, 5),
    encoder_dims: Sequence[int] = (96, 192, 384),
    token_local_channels: Sequence[int] = (128, 256),
    token_global_channels: Sequence[int] = (512,),
    local_radius: Sequence[float] = (0.32, 0.64, 1.28),
    decoder_depths: Sequence[int] = (1, 1),
    decoder_dims: Sequence[int] = (384, 192),
    decoder_up_blocks: Sequence[int] = (1, 1),
    num_heads: int = 6,
    drop_path: float = 0.1,
)

Bases: PretrainingModel

Implementation of the Point-M2AE pre-training model.

Point-M2AE: Multi-scale Masked Autoencoders for Hierarchical Point Cloud Pre-training. This implementation is adapted from the official repository ZrrSkywalker/Point-M2AE.

The model masks tokens at the coarsest stage, back-propagates the mask through the multi-scale grouping, encodes the visible tokens with the hierarchical encoder and reconstructs the masked local neighborhoods with a hierarchical decoder and a reconstruction head.

Parameters:

  • in_channels (int) –

    The number of input channels (unused; coordinates drive the grouping).

  • group_sizes (Sequence[int], default: (16, 8, 8) ) –

    The neighborhood size per stage.

  • num_groups (Sequence[int], default: (512, 256, 64) ) –

    The number of centers per stage.

  • mask_ratio (float, default: 0.8 ) –

    The fraction of coarsest-stage tokens to mask.

  • encoder_depths (Sequence[int], default: (5, 5, 5) ) –

    The number of encoder blocks per stage.

  • encoder_dims (Sequence[int], default: (96, 192, 384) ) –

    The encoder channel width per stage.

  • token_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the first-stage token embedder's per-point MLP.

  • token_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the first-stage token embedder's per-group MLP.

  • local_radius (Sequence[float], default: (0.32, 0.64, 1.28) ) –

    The local-attention radius per stage (disabled during pre-training).

  • decoder_depths (Sequence[int], default: (1, 1) ) –

    The number of decoder blocks per stage.

  • decoder_dims (Sequence[int], default: (384, 192) ) –

    The decoder channel width per stage.

  • decoder_up_blocks (Sequence[int], default: (1, 1) ) –

    The number of residual blocks in each feature-propagation stage.

  • num_heads (int, default: 6 ) –

    The number of attention heads.

  • drop_path (float, default: 0.1 ) –

    The maximum stochastic depth rate.

Shape
  • pos: \((N, 3)\)
  • batch: \((N,)\)
  • Output: a tuple (pred, target, pred_batch, target_batch): reconstructed and ground-truth neighborhood points, \((L \cdot k_0, 3)\) and \((L \cdot k_1, 3)\), each packed with the neighborhood index of its points.

Methods:

  • configure_h_encoder –

    Build the hierarchical encoder with multi-scale back-propagated masking.

  • configure_h_decoder –

    Build the per-stage decoder transformer blocks with a linearly scaled stochastic-depth schedule.

  • configure_decoder_pos_embeds –

    Build the per-stage decoder positional-embedding MLPs.

  • configure_token_prop –

    Build the feature-propagation modules upsampling decoder tokens between consecutive stages.

  • configure_rec_head –

    Build the linear head predicting the coordinates of each masked finest-stage neighborhood.

configure_h_encoder [source] ΒΆ

configure_h_encoder() -> HierarchicalEncoderMAE

Build the hierarchical encoder with multi-scale back-propagated masking.

configure_h_decoder [source] ΒΆ

configure_h_decoder() -> ModuleList

Build the per-stage decoder transformer blocks with a linearly scaled stochastic-depth schedule.

configure_decoder_pos_embeds [source] ΒΆ

configure_decoder_pos_embeds() -> ModuleList

Build the per-stage decoder positional-embedding MLPs.

configure_token_prop [source] ΒΆ

configure_token_prop() -> ModuleList

Build the feature-propagation modules upsampling decoder tokens between consecutive stages.

configure_rec_head [source] ΒΆ

configure_rec_head() -> Linear

Build the linear head predicting the coordinates of each masked finest-stage neighborhood.

PointMAEClassification [source] ΒΆ

PointMAEClassification(
    in_channels: int,
    num_classes: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    drop_path: float = 0.1,
    dropout: float = 0.5,
    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,
    spatial_dim: int = 3,
)

Bases: ClassificationModel

Point-MAE classification model, as in Masked Autoencoders for Point Cloud Self-supervised Learning, adapted from Pang-Yatian/Point-MAE.

Tokenizes local groups with a mini-PointNet encoder, prepends a class token, processes the sequence with a plain transformer encoder, and classifies the concatenation of the class token and the max-pooled patch tokens.

Parameters:

  • in_channels (int) –

    Number of per-point feature channels concatenated to the coordinates (\(0\) for coordinates only).

  • num_classes (int) –

    Number of output classes.

  • embed_dim (int, default: 384 ) –

    Token channels.

  • depth (int, default: 12 ) –

    Number of transformer blocks.

  • num_heads (int, default: 6 ) –

    Number of attention heads.

  • num_groups (int, default: 64 ) –

    Number of groups (FPS centers) per sample.

  • group_size (int, default: 32 ) –

    Number of neighbors per group.

  • encoder_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the patch embedder's per-point MLP.

  • encoder_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the patch embedder's per-group MLP.

  • pos_embed_channels (Sequence[int], default: (128,) ) –

    Hidden widths of the positional-embedding MLP.

  • drop_path (float, default: 0.1 ) –

    Stochastic-depth drop-path rate.

  • dropout (float, default: 0.5 ) –

    Dropout rate in the classification head.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra arguments for the activation.

  • norm (Union[str, Callable, None], default: LayerNorm ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra arguments for the normalization.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: \((B, C)\) where \(B\) is the batch size and \(C\) is num_classes.

Methods:

  • configure_encoder –

    Build the mini-PointNet patch embedder tokenizing each local group.

  • configure_pos_embed –

    Build the positional-embedding MLP mapping group centers to token channels.

  • configure_blocks –

    Build the transformer encoder with a linearly scaled stochastic-depth schedule.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointPatchEmbed

Build the mini-PointNet patch embedder tokenizing each local group.

configure_pos_embed [source] ΒΆ

configure_pos_embed() -> MLP

Build the positional-embedding MLP mapping group centers to token channels.

configure_blocks [source] ΒΆ

configure_blocks() -> TransformerEncoder

Build the transformer encoder with a linearly scaled stochastic-depth schedule.

PointMAEPartSegmentation [source] ΒΆ

PointMAEPartSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    num_categories: int = 16,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 128,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    drop_path: float = 0.1,
    dropout: float = 0.5,
    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,
    spatial_dim: int = 3,
)

Bases: PartSegmentationModel

Point-MAE part-segmentation model, as in Masked Autoencoders for Point Cloud Self-supervised Learning, adapted from Pang-Yatian/Point-MAE.

Tokenizes local groups with a mini-PointNet encoder, processes them with a plain transformer encoder, concatenates the normalized hidden states at three block depths, and propagates the group features back to every point with a PointNet++-style feature propagation (PointNet2FeaturePropagation). A category-conditioned global branch is fused before the per-point classifier. Every sample in the packed batch must contain the same number of points; a ragged batch raises ValueError.

Parameters:

  • in_channels (int) –

    Number of per-point feature channels concatenated to the coordinates (\(0\) for coordinates only).

  • num_classes (int) –

    Number of output part classes (across all categories).

  • num_categories (int, default: 16 ) –

    Number of object categories for the category one-hot branch.

  • embed_dim (int, default: 384 ) –

    Token channels.

  • depth (int, default: 12 ) –

    Number of transformer blocks.

  • num_heads (int, default: 6 ) –

    Number of attention heads.

  • num_groups (int, default: 128 ) –

    Number of groups (FPS centers) per sample.

  • group_size (int, default: 32 ) –

    Number of neighbors per group.

  • encoder_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the patch embedder's per-point MLP.

  • encoder_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the patch embedder's per-group MLP.

  • pos_embed_channels (Sequence[int], default: (128,) ) –

    Hidden widths of the positional-embedding MLP.

  • drop_path (float, default: 0.1 ) –

    Stochastic-depth drop-path rate.

  • dropout (float, default: 0.5 ) –

    Dropout rate in the per-point head.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra arguments for the activation.

  • norm (Union[str, Callable, None], default: LayerNorm ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra arguments for the normalization.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

Shape
  • Input: \((N, 3)\), \((N,)\), and a category one-hot \((B, \text{num\_categories})\).
  • Output: \((N, C)\) logits where \(C\) is num_classes.

Methods:

  • configure_encoder –

    Build the mini-PointNet patch embedder tokenizing each local group.

  • configure_pos_embed –

    Build the positional-embedding MLP mapping group centers to token channels.

  • configure_blocks –

    Build the transformer encoder with a linearly scaled stochastic-depth schedule.

  • configure_label_conv –

    Build the MLP embedding the category one-hot for the global branch.

  • configure_propagation_0 –

    Build the feature-propagation module interpolating group features back to every point.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the per-point features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the per-point features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointPatchEmbed

Build the mini-PointNet patch embedder tokenizing each local group.

configure_pos_embed [source] ΒΆ

configure_pos_embed() -> MLP

Build the positional-embedding MLP mapping group centers to token channels.

configure_blocks [source] ΒΆ

configure_blocks() -> TransformerEncoder

Build the transformer encoder with a linearly scaled stochastic-depth schedule.

configure_label_conv [source] ΒΆ

configure_label_conv() -> MLP

Build the MLP embedding the category one-hot for the global branch.

configure_propagation_0 [source] ΒΆ

configure_propagation_0() -> PointNet2FeaturePropagation

Build the feature-propagation module interpolating group features back to every point.

PointMAEPretraining [source] ΒΆ

PointMAEPretraining(
    in_channels: int,
    *,
    embed_dim: int = 384,
    encoder_depth: int = 12,
    decoder_depth: int = 4,
    num_heads: int = 6,
    decoder_num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    mask_ratio: float = 0.6,
    drop_path: float = 0.1,
    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,
    spatial_dim: int = 3,
)

Bases: PretrainingModel

Point-MAE masked-autoencoder pretraining model, as in Masked Autoencoders for Point Cloud Self-supervised Learning, adapted from Pang-Yatian/Point-MAE.

Masks a fraction of the group tokens, encodes the visible tokens, then reconstructs the masked groups' centered coordinates from learnable mask tokens with a transformer decoder and a per-token coordinate head. forward returns the predicted and target coordinates of the masked groups, each packed with the group index of its points, for a set-to-set reconstruction objective such as chamfer_distance from torch_pointcloud.losses.

Parameters:

  • in_channels (int) –

    Number of input channels (unused beyond coordinates; kept for the registry contract).

  • embed_dim (int, default: 384 ) –

    Token channels.

  • encoder_depth (int, default: 12 ) –

    Number of encoder transformer blocks.

  • decoder_depth (int, default: 4 ) –

    Number of decoder transformer blocks.

  • num_heads (int, default: 6 ) –

    Number of encoder attention heads.

  • decoder_num_heads (int, default: 6 ) –

    Number of decoder attention heads.

  • num_groups (int, default: 64 ) –

    Number of groups (FPS centers) per sample.

  • group_size (int, default: 32 ) –

    Number of neighbors per group.

  • encoder_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the patch embedder's per-point MLP.

  • encoder_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the patch embedder's per-group MLP.

  • pos_embed_channels (Sequence[int], default: (128,) ) –

    Hidden widths of the encoder and decoder positional-embedding MLPs.

  • mask_ratio (float, default: 0.6 ) –

    Fraction of tokens to mask.

  • drop_path (float, default: 0.1 ) –

    Stochastic-depth drop-path rate.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra arguments for the activation.

  • norm (Union[str, Callable, None], default: LayerNorm ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra arguments for the normalization.

  • spatial_dim (int, default: 3 ) –

    Spatial dimension of the input point cloud.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: predicted and target groups, each of shape \((B \cdot M_\text{mask}, M, 3)\).

Methods:

configure_MAE_encoder [source] ΒΆ

configure_MAE_encoder() -> MaskTransformer

Build the masked transformer encoding the visible group tokens.

configure_decoder_pos_embed [source] ΒΆ

configure_decoder_pos_embed() -> MLP

Build the decoder positional-embedding MLP mapping group centers to token channels.

configure_MAE_decoder [source] ΒΆ

configure_MAE_decoder() -> TransformerDecoder

Build the transformer decoder reconstructing the masked tokens.

configure_increase_dim [source] ΒΆ

configure_increase_dim() -> MLP

Build the per-token linear head predicting the coordinates of each masked group.

PointMambaClassification [source] ΒΆ

PointMambaClassification(
    in_channels: int,
    num_classes: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_groups: int = 64,
    group_size: int = 32,
    drop_path: float = 0.1,
    use_cls_token: bool = False,
    spatial_dim: int = 3,
    patch_local_channels: Sequence[int] = (128, 256),
    patch_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    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,
    dropout: float = 0.5,
    global_pool: AdaptivePoolLike = "mean",
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
)

Bases: ClassificationModel

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointMambaEncoder

Build the PointMambaEncoder backbone.

PointMambaPretraining [source] ΒΆ

PointMambaPretraining(
    in_channels: int,
    *,
    embed_dim: int = 384,
    encoder_depth: int = 12,
    decoder_depth: int = 4,
    num_groups: int = 64,
    group_size: int = 32,
    mask_ratio: float = 0.6,
    drop_path: float = 0.1,
    spatial_dim: int = 3,
    patch_local_channels: Sequence[int] = (128, 256),
    patch_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    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,
)

Bases: PretrainingModel

PointMamba masked-autoencoder pre-training model, as described in the paper PointMamba: A Simple State Space Model for Point Cloud Analysis. This implementation is adapted from the official repository LMD0311/PointMamba.

Masks a fraction of the serialized patch tokens, encodes the visible tokens with Mamba blocks, reconstructs the masked patches' relative coordinates with a Mamba decoder and a linear head, and returns the predicted and target patch coordinates, each packed with the patch index of its points, for a set-to-set reconstruction objective such as chamfer_distance from torch_pointcloud.losses.

Parameters:

  • in_channels (int) –

    The number of input channels.

  • embed_dim (int, default: 384 ) –

    The number of token channels.

  • encoder_depth (int, default: 12 ) –

    The number of encoder Mamba blocks.

  • decoder_depth (int, default: 4 ) –

    The number of decoder Mamba blocks.

  • num_groups (int, default: 64 ) –

    The number of patches to sample.

  • group_size (int, default: 32 ) –

    The number of neighbors to consider for each patch.

  • mask_ratio (float, default: 0.6 ) –

    The fraction of patch tokens to mask.

  • drop_path (float, default: 0.1 ) –

    The maximum stochastic-depth rate, linearly scaled across the Mamba blocks.

  • spatial_dim (int, default: 3 ) –

    The dimension of the spatial features.

  • patch_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the patch embedder's per-point MLP.

  • patch_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the patch embedder's per-patch MLP.

  • pos_embed_channels (Sequence[int], default: (128,) ) –

    Hidden widths of the positional-embedding MLPs.

  • act (Union[str, Callable, None], default: 'relu' ) –

    The activation function to use.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    The keyword arguments to pass to 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' ) –

    The normalization function to use.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    The keyword arguments to pass to the normalization function.

  • bias (Union[bool, List[bool]], default: True ) –

    Whether to use bias in the MLPs.

Shape
  • Input: \((N, C_\text{in})\) features (or None), \((N, 3)\) coordinates, and a \((N,)\) batch index.
  • Output: predicted and target patches, each of shape \((B \cdot M_\text{mask}, \text{group\_size}, 3)\).

Methods:

  • configure_encoder –

    Build the PointMambaEncoderMAE backbone, which encodes only the visible patches.

  • configure_decoder –

    Build the PointMambaDecoderMAE that reconstructs the masked patches.

configure_encoder [source] ΒΆ

configure_encoder() -> PointMambaEncoderMAE

Build the PointMambaEncoderMAE backbone, which encodes only the visible patches.

configure_decoder [source] ΒΆ

configure_decoder() -> PointMambaDecoderMAE

Build the PointMambaDecoderMAE that reconstructs the masked patches.

PointTransformerClassification [source] ΒΆ

PointTransformerClassification(
    in_channels: int,
    num_classes: int,
    *,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_attention_groups: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    ratios: Sequence[float],
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    global_pool: PoolLike = "max",
    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,
)

Bases: ClassificationModel

Point Transformer classification model from the paper Point Transformer by Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, Vladlen Koltun.

A hierarchical encoder of vector-attention PointTransformerBlock stages interleaved with PointTransformerTransitionDown downsampling, followed by global pooling and a linear head.

Parameters:

  • in_channels (int) –

    Number of input feature channels. Pass \(0\) to use the raw positions as features.

  • num_classes (int) –

    Number of output classes.

  • encoder_channels (Sequence[int]) –

    Feature width of each encoder stage.

  • encoder_depths (Sequence[int]) –

    Number of PointTransformerBlock blocks per encoder stage.

  • encoder_attention_groups (Sequence[int]) –

    Number of shared-weight vector-attention groups per encoder stage.

  • encoder_num_neighbors (Sequence[int]) –

    Number of neighbors in the \(k\)-NN graph of each encoder stage.

  • ratios (Sequence[float]) –

    Farthest-point-sampling keep ratio for each downsampling transition (length one less than the number of encoder stages).

  • spatial_dim (int, default: 3 ) –

    Dimensionality of the point coordinates.

  • add_self_loops (bool, default: False ) –

    Whether to add self-loops to each neighborhood graph.

  • global_pool (PoolLike, default: 'max' ) –

    Pooling used to aggregate point features into a per-cloud vector.

  • dropout (float, default: 0.0 ) –

    Dropout probability applied before the classification head.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation used across the network.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the activation factory.

  • act_first (bool, default: False ) –

    Whether to apply the activation before the normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization used across the network.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the normalization factory.

Shape
  • x: \((N, \text{in\_channels})\), or None to fall back to pos.
  • pos: \((N, 3)\) point coordinates.
  • batch: \((N,)\) per-point batch index.
  • output: \((B, \text{num\_classes})\) class logits.

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_embeddings [source] ΒΆ

configure_embeddings() -> MLP

Build the embedding MLP lifting the input features to the first encoder width.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerEncoder

Build the PointTransformerEncoder backbone.

PointTransformerSegmentation [source] ΒΆ

PointTransformerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_attention_groups: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_attention_groups: Sequence[int],
    decoder_num_neighbors: Sequence[int],
    ratios: Sequence[float],
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    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,
)

Bases: SemanticSegmentationModel

Point Transformer segmentation model from the paper Point Transformer by Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, Vladlen Koltun.

An encoder-decoder with skip connections: vector-attention PointTransformerBlock stages with PointTransformerTransitionDown downsampling, a PointTransformerGlobalTransition mixing a global mean feature into the coarsest stage, PointTransformerTransitionUp upsampling mirroring the encoder, and a per-point MLP head.

Parameters:

  • in_channels (int) –

    Number of input feature channels. Pass \(0\) to use the raw positions as features.

  • num_classes (int) –

    Number of output classes.

  • encoder_channels (Sequence[int]) –

    Feature width of each encoder stage.

  • encoder_depths (Sequence[int]) –

    Number of PointTransformerBlock blocks per encoder stage.

  • encoder_attention_groups (Sequence[int]) –

    Number of shared-weight vector-attention groups per encoder stage.

  • encoder_num_neighbors (Sequence[int]) –

    Number of neighbors in the \(k\)-NN graph of each encoder stage.

  • decoder_channels (Sequence[int]) –

    Feature width of each decoder stage (the last entry is the head width).

  • decoder_depths (Sequence[int]) –

    Number of PointTransformerBlock blocks per decoder stage.

  • decoder_attention_groups (Sequence[int]) –

    Number of shared-weight vector-attention groups per decoder stage.

  • decoder_num_neighbors (Sequence[int]) –

    Number of neighbors in the \(k\)-NN graph of each decoder stage.

  • ratios (Sequence[float]) –

    Farthest-point-sampling keep ratio for each downsampling transition (length one less than the number of encoder stages).

  • spatial_dim (int, default: 3 ) –

    Dimensionality of the point coordinates.

  • add_self_loops (bool, default: False ) –

    Whether to add self-loops to each neighborhood graph.

  • dropout (float, default: 0.0 ) –

    Dropout probability applied to the per-point features before the head.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation used across the network.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the activation factory.

  • act_first (bool, default: False ) –

    Whether to apply the activation before the normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization used across the network.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the normalization factory.

Shape
  • x: \((N, \text{in\_channels})\), or None to fall back to pos.
  • pos: \((N, 3)\) point coordinates.
  • batch: \((N,)\) per-point batch index.
  • output: \((N, \text{num\_classes})\) per-point class logits.

Methods:

  • configure_embeddings –

    Build the embedding MLP lifting the input features to the first encoder width.

  • configure_encoder –

    Build the PointTransformerEncoder backbone.

  • configure_decoder –

    Build the PointTransformerDecoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_embeddings [source] ΒΆ

configure_embeddings() -> MLP

Build the embedding MLP lifting the input features to the first encoder width.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerEncoder

Build the PointTransformerEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointTransformerDecoder

Build the PointTransformerDecoder upsampling the coarsest features back through the encoder skips.

PointTransformerV2Classification [source] ΒΆ

PointTransformerV2Classification(
    in_channels: int,
    num_classes: int,
    grid_sizes: Sequence[float] = (0.06, 0.12, 0.24, 0.48),
    encoder_depths: Sequence[int] = (1, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        48,
        96,
        192,
        384,
        512,
    ),
    encoder_attention_groups: Sequence[int] = (
        6,
        12,
        24,
        48,
        64,
    ),
    encoder_num_neighbors: Sequence[int] = (
        8,
        16,
        16,
        16,
        16,
    ),
    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,
    qkv_bias: bool = True,
    attn_drop: float = 0.0,
    pe_multiplier: bool = False,
    pe_bias: bool = True,
    drop_path: float = 0.0,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
)

Bases: ClassificationModel

Implementation of the Point Transformer V2 model for classification as described in the paper Point Transformer V2: Grouped Vector Attention and Partition-based Pooling by Xiaoyang Wu, Yixing Lao, Li Jiang, Xihui Liu, Hengshuang Zhao.

Note

This implementation requires pyg-lib to be installed.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes.

  • encoder_depths (Sequence[int], default: (1, 2, 2, 6, 2) ) –

    Number of encoder blocks for each stage.

  • encoder_channels (Sequence[int], default: (48, 96, 192, 384, 512) ) –

    Number of channels for each encoder block.

  • encoder_attention_groups (Sequence[int], default: (6, 12, 24, 48, 64) ) –

    Number of groups for each encoder block.

  • encoder_num_neighbors (Sequence[int], default: (8, 16, 16, 16, 16) ) –

    Number of edge_index for each encoder block.

  • grid_sizes (Sequence[float], default: (0.06, 0.12, 0.24, 0.48) ) –

    Size of the grid for each stage.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer to use.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function to use.

  • qkv_bias (bool, default: True ) –

    Whether to use bias in the QKV linear layer.

  • pe_multiplier (bool, default: False ) –

    Whether to use a multiplier for the PE.

  • pe_bias (bool, default: True ) –

    Whether to use bias in the PE linear layer.

  • drop_path (float, default: 0.0 ) –

    Drop path rate.

  • dropout (float, default: 0.0 ) –

    Dropout rate.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling method to use.

Inputs

x: Float tensor of shape \((N, \text{in\_channels})\). pos: Float tensor of shape \((N, 3)\). batch: Long tensor of shape \((N,)\).

Outputs

Logits tensor of shape \((B, \text{num\_classes})\).

Methods:

  • configure_stem –

    Build the linear stem lifting the input features to the first encoder channel.

  • configure_encoder –

    Build the PointTransformerV2Encoder backbone.

  • reset_classifier –

    Resets the classification head with new parameters.

  • forward_features –

    Forward features through the encoder blocks, before the global pooling.

  • forward_head –

    Forward pass of the classification head from pre-pooling x.

  • forward –

    Forward pass of the classification network.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_stem [source] ΒΆ

configure_stem() -> Module

Build the linear stem lifting the input features to the first encoder channel.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV2Encoder

Build the PointTransformerV2Encoder backbone.

reset_classifier [source] ΒΆ

reset_classifier(
    num_classes: int,
    global_pool: Optional[PoolLike] = None,
    **kwargs: Any,
) -> None

Resets the classification head with new parameters.

Note

To set an empty classification head, use num_classes=0.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • global_pool (Optional[PoolLike], default: None ) –

    Pooling method to aggregate point x ("max" or "mean"). If None, keeps the current pooling.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the classification head.

forward_features [source] ΒΆ

forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: Literal[True],
) -> Tuple[
    Tensor, Tensor, Tensor, List[PooledFeaturesDict]
]
forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: Literal[False] = False,
) -> Tuple[Tensor, Tensor, Tensor]
forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: bool = False,
) -> Any

Forward features through the encoder blocks, before the global pooling.

Parameters:

  • x (OptTensor) –

    Additional point features of shape \((N, \text{features\_dim})\).

  • pos (Tensor) –

    Coordinates of shape \((N, 3)\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • return_intermediates (bool, default: False ) –

    Whether to return the intermediate features.

Returns:

  • x ( Any ) –

    Pre-pooling features of shape \((N, \text{embedding\_dim})\).

  • pos ( Any ) –

    Coordinates of shape \((N, 3)\).

  • batch ( Any ) –

    Batch indices for each point of shape \((N,)\).

  • intermediates ( Any ) –

    If return_intermediates is True, a list of dictionaries containing the intermediate features, coordinates, batch indices and pooling inverse for each encoder block.

forward_head [source] ΒΆ

forward_head(
    x: Tensor, batch: Tensor, pre_logits: bool = False
) -> Tensor

Forward pass of the classification head from pre-pooling x.

Parameters:

  • x (Tensor) –

    Pre-pooling features of shape \((N, \text{embedding\_dim})\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • pre_logits (bool, default: False ) –

    Whether to return pre-logits.

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

forward [source] ΒΆ

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

Forward pass of the classification network.

Parameters:

  • x (OptTensor) –

    Additional point features of shape \((N, \text{features\_dim})\).

  • pos (Tensor) –

    Coordinates of shape \((N, 3)\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

PointTransformerV2Segmentation [source] ΒΆ

PointTransformerV2Segmentation(
    in_channels: int,
    num_classes: int,
    grid_sizes: Sequence[float] = (0.06, 0.12, 0.24, 0.48),
    encoder_depths: Sequence[int] = (1, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        48,
        96,
        192,
        384,
        512,
    ),
    encoder_attention_groups: Sequence[int] = (
        6,
        12,
        24,
        48,
        64,
    ),
    encoder_num_neighbors: Sequence[int] = (
        8,
        16,
        16,
        16,
        16,
    ),
    decoder_depths: Sequence[int] = (1, 1, 1, 1),
    decoder_channels: Sequence[int] = (384, 192, 96, 48),
    decoder_attention_groups: Sequence[int] = (
        48,
        24,
        12,
        6,
    ),
    decoder_num_neighbors: Sequence[int] = (16, 16, 16, 16),
    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,
    qkv_bias: bool = True,
    attn_drop: float = 0.0,
    pe_multiplier: bool = False,
    pe_bias: bool = True,
    drop_path: float = 0.0,
    dropout: float = 0.0,
)

Bases: SemanticSegmentationModel

Implementation of the Point Transformer V2 model for semantic segmentation as described in the paper Point Transformer V2: Grouped Vector Attention and Partition-based Pooling by Xiaoyang Wu, Yixing Lao, Li Jiang, Xihui Liu, Hengshuang Zhao.

Note

This implementation requires pyg-lib to be installed.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes.

  • encoder_depths (Sequence[int], default: (1, 2, 2, 6, 2) ) –

    Number of encoder blocks for each stage.

  • encoder_channels (Sequence[int], default: (48, 96, 192, 384, 512) ) –

    Number of channels for each encoder block.

  • encoder_attention_groups (Sequence[int], default: (6, 12, 24, 48, 64) ) –

    Number of groups for each encoder block.

  • encoder_num_neighbors (Sequence[int], default: (8, 16, 16, 16, 16) ) –

    Number of edge_index for each encoder block.

  • decoder_depths (Sequence[int], default: (1, 1, 1, 1) ) –

    Number of decoder blocks per stage.

  • decoder_channels (Sequence[int], default: (384, 192, 96, 48) ) –

    Number of channels for each decoder block.

  • decoder_attention_groups (Sequence[int], default: (48, 24, 12, 6) ) –

    Number of groups for each decoder block.

  • decoder_num_neighbors (Sequence[int], default: (16, 16, 16, 16) ) –

    Neighbor count for each decoder block.

  • grid_sizes (Sequence[float], default: (0.06, 0.12, 0.24, 0.48) ) –

    Size of the grid for each stage.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer to use.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function to use.

  • qkv_bias (bool, default: True ) –

    Whether to use bias in the QKV linear layer.

  • pe_multiplier (bool, default: False ) –

    Whether to use a multiplier for the PE.

  • pe_bias (bool, default: True ) –

    Whether to use bias in the PE linear layer.

  • drop_path (float, default: 0.0 ) –

    Drop path rate.

  • dropout (float, default: 0.0 ) –

    Dropout rate.

  • attn_drop (float, default: 0.0 ) –

    Attention dropout rate.

Inputs

x: Float tensor of shape \((N, \text{in\_channels})\). pos: Float tensor of shape \((N, 3)\). batch: Long tensor of shape \((N,)\).

Outputs

Segmentation logits of shape \((N, \text{num\_classes})\).

Methods:

  • configure_stem –

    Build the linear stem lifting the input features to the first encoder channel.

  • configure_encoder –

    Build the PointTransformerV2Encoder backbone.

  • configure_decoder –

    Build the PointTransformerV2Decoder upsampling the coarsest features back through the encoder skips.

  • reset_classifier –

    Resets the head with new class parameters.

  • forward_features –

    Forward features through the encoder blocks, before the global pooling.

  • forward_head –

    Forward pass of the classification head from up-sampled x.

  • forward –

    Forward pass of the model.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the per-point decoder features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the per-point decoder features entering the head.

configure_stem [source] ΒΆ

configure_stem() -> Module

Build the linear stem lifting the input features to the first encoder channel.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV2Encoder

Build the PointTransformerV2Encoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointTransformerV2Decoder

Build the PointTransformerV2Decoder upsampling the coarsest features back through the encoder skips.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int, **kwargs: Any) -> None

Resets the head with new class parameters.

Note

To set an empty head, use num_classes=0.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the segmentation head.

forward_features [source] ΒΆ

forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: Literal[True],
) -> Tuple[
    Tensor, Tensor, Tensor, List[PooledFeaturesDict]
]
forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: Literal[False] = False,
) -> Tuple[Tensor, Tensor, Tensor]
forward_features(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    return_intermediates: bool = False,
) -> Any

Forward features through the encoder blocks, before the global pooling.

Parameters:

  • x (OptTensor) –

    Additional point features of shape \((N, \text{features\_dim})\).

  • pos (Tensor) –

    Coordinates of shape \((N, 3)\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • return_intermediates (bool, default: False ) –

    Whether to return the intermediate features.

Returns:

  • x ( Any ) –

    Pre-pooling features of shape \((N, \text{embedding\_dim})\).

  • pos ( Any ) –

    Coordinates of shape \((N, 3)\).

  • batch ( Any ) –

    Batch indices for each point of shape \((N,)\).

  • intermediates ( Any ) –

    If return_intermediates is True, a list of dictionaries containing the intermediate features, coordinates, batch indices and pooling inverse for each encoder block.

forward_head [source] ΒΆ

forward_head(x: Tensor, pre_logits: bool = False) -> Tensor

Forward pass of the classification head from up-sampled x.

Parameters:

  • x (Tensor) –

    Pre-pooling features of shape \((N, \text{embedding\_dim})\).

  • pre_logits (bool, default: False ) –

    Whether to return pre-logits.

Returns:

  • Tensor –

    Segmentation logits of shape \((N, \text{num\_classes})\).

forward [source] ΒΆ

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

Forward pass of the model.

Parameters:

  • x (OptTensor) –

    Additional point features of shape \((N, \text{features\_dim})\).

  • pos (Tensor) –

    Coordinates of shape \((N, 3)\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tensor –

    Segmentation logits of shape \((N, \text{num\_classes})\).

PointTransformerV3Classification [source] ΒΆ

PointTransformerV3Classification(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (2, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        512,
    ),
    encoder_num_heads: Sequence[int] = (2, 4, 8, 16, 32),
    encoder_patch_size: Sequence[int] = (
        48,
        48,
        48,
        48,
        48,
    ),
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    rope_base: float = 10.0,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    pooling: str = "serialized",
    stem_type: str = "sparse_conv",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
    pdnorm_conditions: Optional[Sequence[str]] = None,
    condition: Optional[str] = None,
)

Bases: ClassificationModel

PyTorch implementation of the Point Transformer V3 model, as described in the paper Point Transformer V3: Simpler, Faster, Stronger by Xiaoyang Wu, Li Jiang, Peng-Shuai Wang, Zhijian Liu, Xihui Liu, Yu Qiao, Wanli Ouyang, Tong He, Hengshuang Zhao.

This implementation is based on the original implementation from Pointcept.

Important

This model requires spconv to be installed. It is also recommended to install flash-attn for faster attention. The registered configurations construct with use_flash_attn=True, which requires flash-attn and a CUDA device; pass use_flash_attn=False to run without it. The xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run on CPU. In addition, it is recommended to install ocnn if you want to use more serialization orders.

Parameters:

  • in_channels (int) –

    Number of input channels (corresponding to the number of features).

  • num_classes (int) –

    Number of output classes.

  • serialization_orders (Sequence[SerializationOrder], default: ('hilbert', 'hilbert-trans') ) –

    Serialization orders to use for the PointTransformerV3Encoder.

  • shuffle_serialization_orders (bool, default: True ) –

    Whether to shuffle the serialization orders each step.

  • strides (Sequence[int], default: (2, 2, 2, 2) ) –

    Downsampling strides between encoder stages.

  • encoder_depths (Sequence[int], default: (2, 2, 2, 6, 2) ) –

    Number of encoder blocks per stage.

  • encoder_channels (Sequence[int], default: (32, 64, 128, 256, 512) ) –

    Number of channels per stage.

  • encoder_num_heads (Sequence[int], default: (2, 4, 8, 16, 32) ) –

    Number of attention heads per stage.

  • encoder_patch_size (Sequence[int], default: (48, 48, 48, 48, 48) ) –

    Patch size per stage.

  • norm (Union[str, Callable], default: 'batch_norm' ) –

    Normalization layer to use.

  • act (Union[str, Callable], default: 'gelu' ) –

    Activation function to use.

  • mlp_ratio (float, default: 4 ) –

    MLP hidden dimension ratio inside each block.

  • qkv_bias (bool, default: True ) –

    Whether to use bias in the QKV projection.

  • qk_scale (Optional[float], default: None ) –

    Scaling factor for the QK matrix.

  • attn_drop (float, default: 0.0 ) –

    Dropout rate for the attention.

  • proj_drop (float, default: 0.0 ) –

    Dropout rate for the output projection of each block.

  • drop_path (float, default: 0.3 ) –

    Stochastic depth rate.

  • attn_kind (AttentionKind, default: 'default' ) –

    Attention variant: "default", "rpe", or "rope". The "rope" variant requires the real-valued pos argument at forward time.

  • use_flash_attn (bool, default: True ) –

    Whether to use flash attention.

  • upcast_attn (bool, default: False ) –

    Whether to upcast the attention to fp32.

  • upcast_softmax (bool, default: False ) –

    Whether to upcast the softmax in fp32.

  • rope_base (float, default: 10.0 ) –

    RoPE frequency base. Only used when attn_kind="rope".

  • dropout (float, default: 0.0 ) –

    Dropout rate before the classification head.

  • global_pool (PoolLike, default: 'max' ) –

    How to pool point features to a batch-level vector ("max", "mean", etc.).

  • pooling (str, default: 'serialized' ) –

    Pooling between encoder stages ("serialized" or "grid").

  • stem_type (str, default: 'sparse_conv' ) –

    Encoder stem: "sparse_conv" or "linear".

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the activation factory.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the normalization factory.

Inputs

x: Float tensor of shape \((N, \text{in\_channels})\). pos_grid: Int tensor of shape \((N, 3)\) with voxel-grid coordinates. batch: Long tensor of shape \((N,)\). pos: Float tensor of shape \((N, 3)\) with metric coordinates. Required when attn_kind="rope".

Outputs

logits: Float tensor of shape \((N, \text{num\_classes})\).

Methods:

  • configure_encoder –

    Build the PointTransformerV3Encoder backbone.

  • reset_classifier –

    Resets the classification head with new parameters.

  • forward_head –

    Forward pass of the classification head from pre-pooling features.

  • forward –

    Forward pass of the Point Transformer V3 classification network.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

reset_classifier [source] ΒΆ

reset_classifier(
    num_classes: int,
    global_pool: Optional[PoolLike] = None,
    **kwargs: Any,
) -> None

Resets the classification head with new parameters.

Note

To set an empty classification head, use num_classes=0.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • global_pool (Optional[PoolLike], default: None ) –

    Pooling method to aggregate point features ("max" or "mean"). None keeps the current pooling.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the classification head.

forward_head [source] ΒΆ

forward_head(
    x: Tensor, batch: Tensor, pre_logits: bool = False
) -> Tensor

Forward pass of the classification head from pre-pooling features.

Parameters:

  • x (Tensor) –

    Pre-pooling features of shape \((N, \text{embedding\_dim})\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • pre_logits (bool, default: False ) –

    Whether to return pre-logits. Defaults to False.

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

forward [source] ΒΆ

forward(
    x: OptTensor,
    pos_grid: Tensor,
    batch: Tensor,
    pos: OptTensor = None,
    *,
    condition: Union[str, Sequence[str], None] = None,
) -> Tensor

Forward pass of the Point Transformer V3 classification network.

Parameters:

  • x (OptTensor) –

    Additional point features of shape \((N, C)\).

  • pos_grid (Tensor) –

    Integer grid coordinates of shape \((N, 3)\). The encoder uses these to derive the Z-order / Hilbert serialization index, so they must be voxel indices, not float positions.

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • pos (OptTensor, default: None ) –

    Real-valued metric positions of shape \((N, 3)\). Required when attn_kind="rope"; ignored otherwise.

  • condition (Union[str, Sequence[str], None], default: None ) –

    Optional per-batch condition selecting the PDNorm inner norms (keyword-only).

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

PointTransformerV3Segmentation [source] ΒΆ

PointTransformerV3Segmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "hilbert",
        "hilbert-trans",
    ),
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (2, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        512,
    ),
    encoder_num_heads: Sequence[int] = (2, 4, 8, 16, 32),
    encoder_patch_size: Sequence[int] = (
        48,
        48,
        48,
        48,
        48,
    ),
    decoder_depths: Sequence[int] = (2, 2, 2, 2),
    decoder_channels: Sequence[int] = (256, 128, 64, 64),
    decoder_num_heads: Sequence[int] = (16, 8, 4, 4),
    decoder_patch_size: Sequence[int] = (48, 48, 48, 48),
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    shuffle_serialization_orders: bool = True,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    rope_base: float = 10.0,
    dropout: float = 0.0,
    pooling: str = "serialized",
    stem_type: str = "sparse_conv",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
    pdnorm_conditions: Optional[Sequence[str]] = None,
    condition: Optional[str] = None,
)

Bases: SemanticSegmentationModel

PyTorch implementation of the Point Transformer V3 model for segmentation tasks.

Based on the paper Point Transformer V3: Simpler, Faster, Stronger by Xiaoyang Wu, Li Jiang, Peng-Shuai Wang, Zhijian Liu, Xihui Liu, Yu Qiao, Wanli Ouyang, Tong He, Hengshuang Zhao.

This segmentation variant uses an encoder-decoder architecture with skip connections.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of output classes for segmentation.

  • serialization_orders (Sequence[SerializationOrder], default: ('hilbert', 'hilbert-trans') ) –

    Serialization orders to use for the encoder.

  • strides (Sequence[int], default: (2, 2, 2, 2) ) –

    Strides for the downsampling operations.

  • encoder_depths (Sequence[int], default: (2, 2, 2, 6, 2) ) –

    Number of encoder blocks for each stage.

  • encoder_channels (Sequence[int], default: (32, 64, 128, 256, 512) ) –

    Number of channels for each encoder block.

  • encoder_num_heads (Sequence[int], default: (2, 4, 8, 16, 32) ) –

    Number of attention heads for each encoder block.

  • encoder_patch_size (Sequence[int], default: (48, 48, 48, 48, 48) ) –

    Patch size for each encoder block.

  • decoder_depths (Sequence[int], default: (2, 2, 2, 2) ) –

    Number of decoder blocks for each stage.

  • decoder_channels (Sequence[int], default: (256, 128, 64, 64) ) –

    Number of channels for each decoder block.

  • decoder_num_heads (Sequence[int], default: (16, 8, 4, 4) ) –

    Number of attention heads for each decoder block.

  • decoder_patch_size (Sequence[int], default: (48, 48, 48, 48) ) –

    Patch size for each decoder block.

  • norm (Union[str, Callable], default: 'batch_norm' ) –

    Normalization layer to use.

  • act (Union[str, Callable], default: 'gelu' ) –

    Activation function to use.

  • mlp_ratio (float, default: 4 ) –

    Ratio of the hidden dimension to the input dimension.

  • qkv_bias (bool, default: True ) –

    Whether to use bias in the QKV projection.

  • qk_scale (Optional[float], default: None ) –

    Scaling factor for the QK matrix.

  • attn_drop (float, default: 0.0 ) –

    Dropout rate for the attention.

  • proj_drop (float, default: 0.0 ) –

    Dropout rate for the projection.

  • drop_path (float, default: 0.3 ) –

    Dropout rate for the drop path.

  • shuffle_serialization_orders (bool, default: True ) –

    Whether to shuffle the serialization orders.

  • attn_kind (AttentionKind, default: 'default' ) –

    Attention variant: "default", "rpe", or "rope". The "rope" variant requires the real-valued pos argument at forward time.

  • rope_base (float, default: 10.0 ) –

    RoPE frequency base. Only used when attn_kind="rope".

  • use_flash_attn (bool, default: True ) –

    Whether to use flash attention. The registered configurations construct with use_flash_attn=True, which requires flash-attn and a CUDA device; pass use_flash_attn=False to run without it (the xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run on CPU).

  • upcast_attn (bool, default: False ) –

    Whether to upcast the attention.

  • upcast_softmax (bool, default: False ) –

    Whether to upcast the softmax.

  • dropout (float, default: 0.0 ) –

    Dropout on the per-point logits.

  • pooling (str, default: 'serialized' ) –

    Inter-stage pooling ("serialized" or "grid").

  • stem_type (str, default: 'sparse_conv' ) –

    Encoder stem ("sparse_conv" or "linear").

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the activation factory.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the normalization factory.

Methods:

  • configure_encoder –

    Build the PointTransformerV3Encoder backbone.

  • configure_decoder –

    Build the PointTransformerV3Decoder upsampling the coarsest features back through the encoder skips.

  • reset_classifier –

    Resets the segmentation head with new parameters.

  • forward –

    Forward pass of the Point Transformer V3 segmentation network.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the per-point decoder features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the per-point decoder features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointTransformerV3Decoder

Build the PointTransformerV3Decoder upsampling the coarsest features back through the encoder skips.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int, **kwargs: Any) -> None

Resets the segmentation head with new parameters.

Note

To set an empty segmentation head, use num_classes=0.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the segmentation head.

forward [source] ΒΆ

forward(
    x: Tensor,
    pos_grid: Tensor,
    batch: Tensor,
    pos: OptTensor = None,
    *,
    condition: Union[str, Sequence[str], None] = None,
) -> Tensor

Forward pass of the Point Transformer V3 segmentation network.

Parameters:

  • x (Tensor) –

    Per-point features of shape \((N, C)\).

  • pos_grid (Tensor) –

    Integer grid coordinates of shape \((N, 3)\) used for serialization.

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • pos (OptTensor, default: None ) –

    Real-valued metric positions of shape \((N, 3)\). Required when attn_kind="rope"; ignored otherwise.

  • condition (Union[str, Sequence[str], None], default: None ) –

    Optional per-batch condition selecting the PDNorm inner norms (keyword-only).

Returns:

  • Tensor –

    Per-point segmentation logits of shape \((N, \text{num\_classes})\).

PointCNNClassification [source] ΒΆ

PointCNNClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    channels: Sequence[int],
    kernel_sizes: Sequence[int],
    ratios: Sequence[float],
    hidden_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    dilations: Sequence[int] = (1, 1, 1, 1),
    bias: bool = True,
    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,
    dropout: float = 0.0,
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    global_pool: PoolLike = "max",
)

Bases: ClassificationModel

Classification model 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 classification model consists of a encoder of XConv layers and FPS downsampling layers, and a MLP classification head.

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the features entering the head.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> Module

Build the PointCNNEncoder backbone.

PointCNNSegmentation [source] ΒΆ

PointCNNSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    channels: Sequence[int],
    hidden_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    kernel_sizes: Sequence[int],
    dilations: Sequence[int] = (1, 1, 1, 1),
    ratios: Sequence[float],
    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,
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
)

Bases: SemanticSegmentationModel

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

An encoder of XConv layers with FPS downsampling, a decoder of XConv layers upsampling back to the skip resolutions, and a per-point MLP head.

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PointCNNEncoder

Build the PointCNNEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointCNNDecoder

Build the PointCNNDecoder, mirroring the decimating encoder stages in reverse.

PointConvDensityClassification [source] ΒΆ

PointConvDensityClassification(
    in_channels: int,
    num_classes: int,
    *,
    channels: Sequence[Sequence[int]] = (
        [64, 64, 128],
        [128, 128, 256],
        [256, 512, 1024],
    ),
    num_neighbors: Sequence[int] = (32, 64, 1024),
    bandwidths: Sequence[float] = (0.1, 0.2, 0.4),
    ratios: Sequence[float] = (0.5, 0.25, 0.0),
    density_channels: Sequence[int] = (16, 8),
    weight_channels: Sequence[int] = (8, 8),
    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.5,
    global_pool: PoolLike = "mean",
    head_channels: Sequence[int] = (512, 256),
)

Bases: ClassificationModel

PointConv classification model with density re-weighting from PointConv: Deep Convolutional Networks on 3D Point Clouds by Wenxuan Wu, Zhongang Qi, Li Fuxin.

Each set-abstraction stage samples centroids with farthest point sampling, estimates a kernel density per point, and applies a density-reweighted continuous convolution over each \(k\)-NN neighborhood. The last stage pools globally, so the classification head operates on one feature vector per sample.

Shape
  • Input: features of shape \((N, \text{in\_channels})\) (optional), points of shape \((N, \text{spatial\_dim})\), batch indices of shape \((N,)\).
  • Output: logits of shape \((B, \text{num\_classes})\).

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PointConvDensityEncoder

Build the PointConvDensityEncoder backbone.

PointGPTClassification [source] ΒΆ

PointGPTClassification(
    in_channels: int,
    num_classes: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    decoder_depth: int = 4,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    dropout: float = 0.5,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    head_act: Union[str, Callable, None] = "relu",
    global_pool: AdaptivePoolLike = "max",
    spatial_dim: int = 3,
)

Bases: ClassificationModel

PointGPT classification model.

Implements the finetuning model (PointTransformer) of PointGPT: Auto-regressively Generative Pre-training from Point Clouds, adapted from CGuangyan-BIT/PointGPT. Patches are tokenized with a mini-PointNet, ordered by a greedy nearest-neighbor ("simplified Morton") traversal of the centers, and consumed by a causally-masked GPT extractor that prepends a start-of-sequence and a class token. The global feature concatenates the class-token output with the pooled patch outputs (global_pool, max-pool by default), so the head input dimension is \(2 \cdot \text{embed\_dim}\).

Parameters:

  • in_channels (int) –

    The number of input channels (PointGPT uses coordinates only, so \(0\)).

  • num_classes (int) –

    The number of output classes.

  • embed_dim (int, default: 384 ) –

    The transformer / token-embedding dimension.

  • depth (int, default: 12 ) –

    The number of extractor blocks.

  • num_heads (int, default: 6 ) –

    The number of attention heads.

  • num_groups (int, default: 64 ) –

    The number of patches \(G\).

  • group_size (int, default: 32 ) –

    The number of points \(M\) per patch.

  • decoder_depth (int, default: 4 ) –

    The number of generator blocks (kept for weight compatibility).

  • encoder_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the patch embedder's per-point MLP.

  • encoder_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the patch embedder's per-patch MLP.

  • dropout (float, default: 0.5 ) –

    The dropout rate of the classification head.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    The activation of the transformer MLPs.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the activation.

  • head_act (Union[str, Callable, None], default: 'relu' ) –

    The activation of the classification head.

  • global_pool (AdaptivePoolLike, default: 'max' ) –

    The pooling over the patch tokens for the global feature ("max" or "mean").

  • spatial_dim (int, default: 3 ) –

    The spatial dimension of the input point cloud.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: \((B, \text{num\_classes})\).
Example
import torch
from torch_pointcloud.models.pointgpt import PointGPTClassification

model = PointGPTClassification(in_channels=0, num_classes=40)
pos = torch.randn(2048, 3)
batch = torch.cat([torch.zeros(1024), torch.ones(1024)]).long()
logits = model(None, pos, batch)
print(logits.shape)

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> PointPatchEmbed

Build the mini-PointNet patch embedder tokenizing each patch.

configure_pos_embed [source] ΒΆ

configure_pos_embed() -> PositionEmbeddingSine

Build the sinusoidal positional embedding of the patch centers.

configure_blocks [source] ΒΆ

configure_blocks() -> PointGPTExtractor

Build the causally-masked GPT extractor.

configure_generator_blocks [source] ΒΆ

configure_generator_blocks() -> PointGPTGenerator

Build the GPT generator predicting the next patch (unused at finetuning, kept for weight compatibility).

PointGPTPretraining [source] ΒΆ

PointGPTPretraining(
    in_channels: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    decoder_depth: int = 4,
    num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    mask_ratio: float = 0.7,
    keep_attend: int = 10,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    spatial_dim: int = 3,
)

Bases: PretrainingModel

PointGPT auto-regressive generative pretraining model.

Implements the pretraining model (PointGPT / GPT_Transformer) of PointGPT: Auto-regressively Generative Pre-training from Point Clouds, adapted from CGuangyan-BIT/PointGPT.

Patches are tokenized, ordered by a greedy nearest-neighbor ("simplified Morton") traversal, and fed to a causally-masked GPT extractor with an additional column mask that randomly hides patches beyond the first keep_attend tokens (the dual-masking strategy). The generator then predicts the next patch from the extractor features and a relative positional embedding. forward returns the predicted and target patch coordinates, each packed with the patch index of its points, for a set-to-set reconstruction objective such as chamfer_distance from torch_pointcloud.losses.

Parameters:

  • in_channels (int) –

    The number of input channels (\(0\), coordinates only).

  • embed_dim (int, default: 384 ) –

    The transformer / token-embedding dimension.

  • depth (int, default: 12 ) –

    The number of extractor blocks.

  • decoder_depth (int, default: 4 ) –

    The number of generator blocks.

  • num_heads (int, default: 6 ) –

    The number of attention heads (shared by the extractor and the generator).

  • num_groups (int, default: 64 ) –

    The number of patches \(G\).

  • group_size (int, default: 32 ) –

    The number of points \(M\) per patch.

  • encoder_local_channels (Sequence[int], default: (128, 256) ) –

    Hidden widths of the patch embedder's per-point MLP.

  • encoder_global_channels (Sequence[int], default: (512,) ) –

    Hidden widths of the patch embedder's per-patch MLP.

  • mask_ratio (float, default: 0.7 ) –

    The fraction of maskable patches hidden by the column mask.

  • keep_attend (int, default: 10 ) –

    The number of leading patches never hidden by the column mask.

  • act (Union[str, Callable, None], default: 'gelu' ) –

    The activation of the transformer MLPs.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the activation.

  • spatial_dim (int, default: 3 ) –

    The spatial dimension of the input point cloud.

Shape
  • Input: \((N, 3)\) and \((N,)\).
  • Output: \((B \cdot G \cdot M, 3)\) predicted and target coordinates, each with its \((B \cdot G \cdot M,)\) patch index.
Example
import torch
from torch_pointcloud.models.pointgpt import PointGPTPretraining

model = PointGPTPretraining(in_channels=0)
pos = torch.randn(2048, 3)
batch = torch.cat([torch.zeros(1024), torch.ones(1024)]).long()
pred, target, pred_batch, target_batch = model(None, pos, batch)
print(pred.shape, target.shape)

Methods:

configure_encoder [source] ΒΆ

configure_encoder() -> PointPatchEmbed

Build the mini-PointNet patch embedder tokenizing each patch.

configure_pos_embed [source] ΒΆ

configure_pos_embed() -> PositionEmbeddingSine

Build the sinusoidal positional embedding of the patch centers and relative directions.

configure_blocks [source] ΒΆ

configure_blocks() -> PointGPTExtractor

Build the causally-masked GPT extractor.

configure_generator_blocks [source] ΒΆ

configure_generator_blocks() -> PointGPTGenerator

Build the GPT generator predicting the next patch from the extractor features.

PointMLPClassification [source] ΒΆ

PointMLPClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    encoder_channels: Sequence[int],
    num_neighbors: Union[int, Sequence[int]],
    ratios: Union[float, Sequence[float]],
    num_pre_blocks: Union[int, Sequence[int]] = 2,
    num_pos_blocks: Union[int, Sequence[int]] = 2,
    normalize: Literal["center", "anchor"] = "center",
    std_mode: Literal["graph", "batch"] = "graph",
    res_expansion: float = 1.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 = True,
    add_self_loops: bool = False,
    use_pos: bool = True,
    dropout: float = 0.0,
    head_channels: Optional[Sequence[int]] = None,
    head_dropout: float = 0.0,
    global_pool: PoolLike = "max",
)

Bases: ClassificationModel

PointMLP classification model 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.

A pure residual MLP network: each encoder stage samples centroids with farthest point sampling, normalizes each \(k\)-NN neighborhood with a geometric affine module, and applies residual MLP blocks before and after the neighborhood aggregation. Point features are pooled globally after the encoder for classification.

Methods:

  • configure_stem –

    Build the linear stem lifting the input features to the first encoder channel.

  • configure_encoder –

    Build the PointMLPEncoder backbone.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_stem [source] ΒΆ

configure_stem() -> Module

Build the linear stem lifting the input features to the first encoder channel.

configure_encoder [source] ΒΆ

configure_encoder() -> PointMLPEncoder

Build the PointMLPEncoder backbone.

PointMLPSegmentation [source] ΒΆ

PointMLPSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    encoder_channels: Sequence[int],
    num_neighbors: Union[int, Sequence[int]],
    ratios: Union[float, Sequence[float]],
    num_pre_blocks: Union[int, Sequence[int]] = 2,
    num_pos_blocks: Union[int, Sequence[int]] = 2,
    decoder_channels: Sequence[int],
    decoder_blocks: Sequence[Module],
    normalize: Literal["center", "anchor"] = "center",
    std_mode: Literal["graph", "batch"] = "graph",
    res_expansion: float = 1.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 = True,
    add_self_loops: bool = False,
    use_pos: bool = True,
    dropout: float = 0.0,
)

Bases: SemanticSegmentationModel

PointMLP segmentation model 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.

The PointMLP encoder is followed by a decoder of residual feature-propagation blocks with skip connections and a per-point linear head.

Methods:

  • configure_stem –

    Build the linear stem lifting the input features to the first encoder channel.

  • configure_encoder –

    Build the PointMLPEncoder backbone.

  • configure_decoder –

    Build the PointMLPDecoder, mirroring the encoder channels in reverse.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_stem [source] ΒΆ

configure_stem() -> Module

Build the linear stem lifting the input features to the first encoder channel.

configure_encoder [source] ΒΆ

configure_encoder() -> PointMLPEncoder

Build the PointMLPEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointMLPDecoder

Build the PointMLPDecoder, mirroring the encoder channels in reverse.

PointNetClassification [source] ΒΆ

PointNetClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    mlp1_channels: Sequence[int] = (64,),
    mlp2_channels: Sequence[int] = (128, 1024),
    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,
    use_features_transform: bool = True,
    tnet_mlp1_dims: Sequence[int] = (64, 128, 1024),
    tnet_mlp2_dims: Sequence[int] = (512, 256),
    tnet_act: Union[str, Callable, None] = "relu",
    tnet_act_kwargs: Optional[Dict[str, Any]] = None,
    tnet_norm: Union[str, Callable, None] = "batch_norm",
    tnet_norm_kwargs: Optional[Dict[str, Any]] = None,
    head_channels: Sequence[int] = (512, 256),
)

Bases: ClassificationModel

PointNet architecture for 3D point cloud classification tasks as described in the original PointNet paper PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation.

This model implements the complete PointNet classification network as described in the original paper. It consists of a PointNet encoder to extract global features from point clouds, followed by a classification head to predict class probabilities.

Abstract

This implementation follows the official PointNet architecture for classification, achieving invariance to point permutation through max pooling and robustness to geometric transformations through the T-Net modules.

Tip

To set an empty classification head, set num_classes=0.

Tip

You can control the activations, normalization, and dropout rate of the encoder and head. To skip them, set them to None.

Parameters:

  • in_channels (int) –

    Dimension of additional point features.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Dimension of point coordinates.

  • dropout (float, default: 0.0 ) –

    Dropout rate applied within the classification head.

  • global_pool (PoolLike, default: 'max' ) –

    Pooling method to aggregate point features ("max" or "mean").

  • mlp1_channels (Sequence[int], default: (64,) ) –

    Dimensions of encoder's first MLP.

  • mlp2_channels (Sequence[int], default: (128, 1024) ) –

    Dimensions of encoder's second MLP.

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

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization to use.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization layers.

  • use_features_transform (bool, default: True ) –

    Whether to use feature transformation.

  • tnet_mlp1_dims (Sequence[int], default: (64, 128, 1024) ) –

    Dimensions of T-Net first MLP.

  • tnet_mlp2_dims (Sequence[int], default: (512, 256) ) –

    Dimensions of T-Net second MLP.

  • tnet_act (Union[str, Callable, None], default: 'relu' ) –

    Activation function for T-Net.

  • tnet_act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the T-Net activation function.

  • tnet_norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization for T-Net.

  • tnet_norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the T-Net normalization layers.

  • head_channels (Sequence[int], default: (512, 256) ) –

    Hidden dimensions of the classification head MLP.

Shape
  • Input: features of shape \((N, \text{in\_channels})\) (optional), points of shape \((N, \text{spatial\_dim})\)
  • Output: logits of shape \((B, \text{num\_classes})\)

Methods:

  • configure_encoder –

    Build the PointNetEncoder backbone.

  • reset_classifier –

    Resets the classification head with new parameters.

  • forward_features –

    Forward pass of the PointNet encoder, returning pre-pooling features.

  • forward_head –

    Forward pass of the classification head from pre-pooling features.

  • forward –

    Forward pass of the PointNet classification network.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNetEncoder

Build the PointNetEncoder backbone.

reset_classifier [source] ΒΆ

reset_classifier(
    num_classes: int,
    global_pool: Optional[PoolLike] = None,
    **kwargs: Any,
) -> None

Resets the classification head with new parameters.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • global_pool (Optional[PoolLike], default: None ) –

    Pooling method to aggregate point features ("max" or "mean"); None keeps the current pooling.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the classification head.

forward_features [source] ΒΆ

forward_features(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

Forward pass of the PointNet encoder, returning pre-pooling features.

Parameters:

  • x (Optional[Tensor]) –

    Additional point features of shape \((N, \text{in\_channels})\).

  • pos (Tensor) –

    Point coordinates of shape \((N, \text{spatial\_dim})\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tensor –

    Pre-pooling features of shape \((N, \text{mlp2\_dims}[-1])\) where \(N\) is the number of points.

forward_head [source] ΒΆ

forward_head(
    x: Tensor, batch: Tensor, pre_logits: bool = False
) -> Tensor

Forward pass of the classification head from pre-pooling features.

Parameters:

  • x (Tensor) –

    Pre-pooling features of shape \((N, \text{mlp2\_dims}[-1])\) where \(N\) is the number of points.

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • pre_logits (bool, default: False ) –

    Whether to return pre-logits. Defaults to False.

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

forward [source] ΒΆ

forward(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

Forward pass of the PointNet classification network.

Parameters:

  • x (Optional[Tensor]) –

    Additional point features of shape \((N, \text{in\_channels})\).

  • pos (Tensor) –

    Point coordinates of shape \((N, \text{spatial\_dim})\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tensor –

    Classification logits of shape \((B, \text{num\_classes})\).

PointNetSegmentation [source] ΒΆ

PointNetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    dropout: float = 0.3,
    mlp1_channels: Sequence[int] = (64,),
    mlp2_channels: Sequence[int] = (128, 1024),
    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,
    global_pool: PoolLike = "max",
    use_features_transform: bool = True,
    tnet_mlp1_dims: Sequence[int] = (64, 128, 1024),
    tnet_mlp2_dims: Sequence[int] = (512, 256),
    tnet_act: Union[str, Callable, None] = "relu",
    tnet_act_kwargs: Optional[Dict[str, Any]] = None,
    tnet_norm: Union[str, Callable, None] = "batch_norm",
    tnet_norm_kwargs: Optional[Dict[str, Any]] = None,
    seg_head_dims: Sequence[int] = (512, 256, 128),
)

Bases: SemanticSegmentationModel

PointNet architecture for point cloud segmentation tasks as described in the original PointNet paper PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation.

This model implements the segmentation variant of PointNet as described in the original paper. It extracts both local point features and global shape features, combines them for each point, and predicts per-point semantic labels. This architecture enables the network to consider both local geometry and global context for segmentation.

Abstract

The key innovation in the segmentation variant is the concatenation of global features with per-point features, allowing each point's classification to be informed by both local geometry and the global shape context. This enables part segmentation that is aware of the overall object structure.

Parameters:

  • in_channels (int) –

    Dimension of additional point features.

  • num_classes (int) –

    Number of segmentation classes.

  • spatial_dim (int, default: 3 ) –

    Dimension of point coordinates.

  • dropout (float, default: 0.3 ) –

    Dropout rate applied within the segmentation head.

  • mlp1_channels (Sequence[int], default: (64,) ) –

    Dimensions of encoder's first MLP.

  • mlp2_channels (Sequence[int], default: (128, 1024) ) –

    Dimensions of encoder's second MLP.

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

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization to use.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization layers.

  • global_pool (PoolLike, default: 'max' ) –

    Pooling method for global features ("max" or "mean").

  • use_features_transform (bool, default: True ) –

    Whether to use feature transformation.

  • tnet_mlp1_dims (Sequence[int], default: (64, 128, 1024) ) –

    Dimensions of T-Net first MLP.

  • tnet_mlp2_dims (Sequence[int], default: (512, 256) ) –

    Dimensions of T-Net second MLP.

  • tnet_act (Union[str, Callable, None], default: 'relu' ) –

    Activation function for T-Net.

  • tnet_act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the T-Net activation function.

  • tnet_norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization for T-Net.

  • tnet_norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the T-Net normalization layers.

  • seg_head_dims (Sequence[int], default: (512, 256, 128) ) –

    Dimensions of segmentation head MLPs.

Shape
  • Input: features of shape \((N, \text{in\_channels})\) (optional), points of shape \((N, \text{spatial\_dim})\)
  • Output: logits of shape \((N, \text{num\_classes})\)

Methods:

  • forward_decoder –

    Decode encoder features back to per-point resolution.

  • configure_encoder –

    Build the PointNetEncoder backbone.

  • reset_classifier –

    Resets the segmentation head with new parameters.

  • forward_features –

    Forward pass of the PointNet encoder, returning pre-pooling features.

  • forward_head –

    Forward pass of the segmentation head from pre-pooling features.

  • forward –

    Forward pass of the PointNet segmentation network.

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: per-point features concatenated with the global feature.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: per-point features concatenated with the global feature.

forward_decoder [source] ΒΆ

forward_decoder(*args: Any, **kwargs: Any) -> Any

Decode encoder features back to per-point resolution.

Canonical signature: forward_decoder(x, ..., intermediates), consuming the output of forward_features and returning per-point features \((N, C)\). Models whose encoder already emits per-point features (DGCNN, PointNet) have no decoder and raise NotImplementedError.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNetEncoder

Build the PointNetEncoder backbone.

reset_classifier [source] ΒΆ

reset_classifier(
    num_classes: int,
    global_pool: Optional[PoolLike] = None,
    **kwargs: Any,
) -> None

Resets the segmentation head with new parameters.

Parameters:

  • num_classes (int) –

    Number of output classes.

  • global_pool (Optional[PoolLike], default: None ) –

    Pooling method for global features ("max" or "mean"); None keeps the current pooling.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the segmentation head.

forward_features [source] ΒΆ

forward_features(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tuple[Tensor, Tensor]

Forward pass of the PointNet encoder, returning pre-pooling features.

Parameters:

  • x (Optional[Tensor]) –

    Additional point features of shape \((N, \text{in\_channels})\).

  • pos (Tensor) –

    Point coordinates of shape \((N, \text{spatial\_dim})\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tuple[Tensor, Tensor] –

    A tuple of (pre_pooling_features, point_features) where pre_pooling_features is of shape \((N, \text{mlp2\_dims}[-1])\) and point_features is of shape \((N, \text{mlp1\_dims}[-1])\).

forward_head [source] ΒΆ

forward_head(
    x: Tensor,
    point_features: Tensor,
    batch: Tensor,
    pre_logits: bool = False,
) -> Tensor

Forward pass of the segmentation head from pre-pooling features.

Parameters:

  • x (Tensor) –

    Pre-pooling features of shape \((N, \text{mlp2\_dims}[-1])\) where \(N\) is the number of points.

  • point_features (Tensor) –

    Point features of shape \((N, \text{mlp1\_dims}[-1])\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

  • pre_logits (bool, default: False ) –

    Whether to return pre-logits. Defaults to False.

Returns:

  • Tensor –

    Per-point segmentation logits of shape \((N, \text{num\_classes})\).

forward [source] ΒΆ

forward(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

Forward pass of the PointNet segmentation network.

Parameters:

  • x (Optional[Tensor]) –

    Additional point features of shape \((N, \text{in\_channels})\).

  • pos (Tensor) –

    Point coordinates of shape \((N, \text{spatial\_dim})\).

  • batch (Tensor) –

    Batch indices for each point of shape \((N,)\).

Returns:

  • Tensor –

    Per-point segmentation logits of shape \((N, \text{num\_classes})\).

PointNet2Classification [source] ΒΆ

PointNet2Classification(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Optional[int] = None,
    sa_channels: Sequence[
        Sequence[Union[int, Sequence[int]]]
    ],
    aggr_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    aggr_use_pos: bool = False,
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    ratios: Sequence[float],
    radii: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    spatial_dim: int = 3,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
    use_pos: bool = True,
    normalize_pos: bool = True,
    sa_aggr: AggrType = "max",
    dropout: Union[float, Sequence[float]] = 0.0,
    global_pool: PoolLike = "max",
)

Bases: ClassificationModel

PointNet++ classification model from the paper PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space by Charles R. Qi, Li Yi, Hao Su, Leonidas J. Guibas.

This network is a hierarchical point cloud classification model. It processes raw point clouds through a PointNet2Encoder (optional stem + SA blocks), an optional aggregation MLP, global pooling, and a classification head.

Parameters:

  • in_channels (int) –

    Number of input channels (features per point).

  • num_classes (int) –

    Number of output classes.

  • stem_channels (Optional[int], default: None ) –

    Optional number of channels for initial linear projection (inside the encoder).

  • sa_channels (Sequence[Sequence[Union[int, Sequence[int]]]]) –

    List of channel configurations for Set Abstraction (SA) blocks. Each element defines the MLP channels for one SA block. For Multi-Scale Grouping (MSG), provide nested lists of channels.

  • aggr_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Channel sizes for the post-encoder aggregation MLP.

  • ratios (Sequence[float]) –

    Sampling ratios for each SA block (between 0 and 1).

  • radii (Sequence[Union[float, Sequence[float]]]) –

    Search radiuses for each SA block's neighborhood. For MSG, provide a list of radii per block.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Max number of neighbors for each SA block. For MSG, provide a list of neighbor counts per block.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of point coordinates (e.g. 3 for 3D, 2 for 2D).

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function type or callable.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional keyword arguments for the activation function.

  • act_first (bool, default: False ) –

    If True, activation is applied before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type or callable.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional keyword arguments for the normalization layer.

  • bias (bool, default: False ) –

    Whether to use bias in linear layers.

  • use_pos (bool, default: True ) –

    Whether to concatenate per-point relative positions to x.

  • sa_aggr (AggrType, default: 'max' ) –

    Aggregation of the grouped features in each SA block.

  • dropout (Union[float, Sequence[float]], default: 0.0 ) –

    Dropout for the classification head: a single rate shared by every hidden layer, or one rate per hidden layer.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling operation.

Methods:

  • configure_encoder –

    Build the PointNet2Encoder backbone.

  • configure_aggr –

    Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output, after the optional aggregation MLP.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output, after the optional aggregation MLP.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNet2Encoder

Build the PointNet2Encoder backbone.

configure_aggr [source] ΒΆ

configure_aggr() -> Optional[MLP]

Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

PointNet2Segmentation [source] ΒΆ

PointNet2Segmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Optional[int] = None,
    sa_channels: Sequence[
        Sequence[Union[int, Sequence[int]]]
    ],
    aggr_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    fp_channels: Sequence[Sequence[int]],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    ratios: Sequence[float],
    radii: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    spatial_dim: int = 3,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
    use_pos: bool = True,
    normalize_pos: bool = True,
    sa_aggr: AggrType = "max",
    dropout: float = 0.0,
    skip_input: bool = True,
    fp_k: Optional[Union[int, Sequence[int]]] = None,
)

Bases: SemanticSegmentationModel

PointNet++ segmentation model from the paper PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space by Charles R. Qi, Li Yi, Hao Su, Leonidas J. Guibas.

This network is a hierarchical point cloud segmentation model built from a PointNet2Encoder (optional stem + SA blocks), an optional aggregation MLP, a PointNet2Decoder (FP blocks with skip connections), and a per-point classification head.

Parameters:

  • in_channels (int) –

    Number of input channels (features per point).

  • num_classes (int) –

    Number of output classes.

  • stem_channels (Optional[int], default: None ) –

    Optional number of channels for initial linear projection (inside the encoder).

  • sa_channels (Sequence[Sequence[Union[int, Sequence[int]]]]) –

    List of channel configurations for Set Abstraction (SA) blocks. Each element defines the MLP channels for one SA block. For Multi-Scale Grouping (MSG), provide nested lists of channels.

  • aggr_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Channel sizes for the post-encoder aggregation MLP.

  • fp_channels (Sequence[Sequence[int]]) –

    List of channel configurations for Feature Propagation (FP) blocks. Each element defines the MLP channels for one FP block.

  • ratios (Sequence[float]) –

    Sampling ratios for each SA block (between 0 and 1).

  • radii (Sequence[Union[float, Sequence[float]]]) –

    Search radiuses for each SA block's neighborhood. For MSG, provide a list of radii per block.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Max number of neighbors for each SA block. For MSG, provide a list of neighbor counts per block.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of point coordinates (e.g. 3 for 3D, 2 for 2D).

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function type or callable.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional keyword arguments for the activation function.

  • act_first (bool, default: False ) –

    If True, activation is applied before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type or callable.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional keyword arguments for the normalization layer.

  • bias (bool, default: False ) –

    Whether to use bias in linear layers.

  • use_pos (bool, default: True ) –

    Whether to concatenate per-point relative positions to x.

  • sa_aggr (AggrType, default: 'max' ) –

    Aggregation of the grouped features in each SA block.

  • dropout (float, default: 0.0 ) –

    Dropout rate for classification head.

Methods:

  • configure_encoder –

    Build the PointNet2Encoder backbone.

  • configure_aggr –

    Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

  • configure_decoder –

    Build the PointNet2Decoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNet2Encoder

Build the PointNet2Encoder backbone.

configure_aggr [source] ΒΆ

configure_aggr() -> Optional[MLP]

Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

configure_decoder [source] ΒΆ

configure_decoder() -> PointNet2Decoder

Build the PointNet2Decoder upsampling the coarsest features back through the encoder skips.

PointNeXtClassification [source] ΒΆ

PointNeXtClassification(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    stem_plain_last: bool = False,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_expansion: Union[int, Sequence[int]] = 4,
    sa_layers: int = 1,
    sa_use_res: bool = True,
    ratios: Sequence[float],
    radiuses: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    add_self_loops: bool = False,
    spatial_dim: int = 3,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: Union[bool, List[bool]] = True,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    head_channels: Optional[Sequence[int]] = None,
    global_sa_channels: Optional[Sequence[int]] = None,
)

Bases: ClassificationModel

PointNeXt classification model as described in the paper PointNeXt: Revisiting PointNet++ with Improved Training and Scaling Strategies by Guocheng Qian, Yuchen Li, Houwen Peng, Jinjie Mai, Hasan Abed Al Kader Hammoud, Mohamed Elhoseiny, Bernard Ghanem.

PointNeXt modernizes PointNet++ through improved training strategies and architectural enhancements, achieving state-of-the-art performance while maintaining efficiency. The model introduces Inverted Residual MLP (InvResMLP) blocks, separable MLPs, and relative position normalization to enable effective model scaling.

Parameters:

  • in_channels (int) –

    Number of input channels (typically 3 for XYZ coordinates, or 6 for XYZ + RGB, or more with additional features like normal).

  • num_classes (int) –

    Number of output classes for classification.

  • stem_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Number of channels in the stem MLP layer(s) that map input to higher dimension. If None, no stem is used.

  • encoder_channels (Sequence[int]) –

    Channel dimensions for each encoder block. The number of channels should be equal to the number of blocks + 1.

  • encoder_depths (Sequence[int]) –

    Number of blocks in each encoder stage after the initial SA block.

  • encoder_expansion (Union[int, Sequence[int]], default: 4 ) –

    Expansion ratio to determine the hidden channels of the encoder blocks.

  • ratios (Sequence[float]) –

    Downsampling sampling ratios for each encoder stage. The number of ratios should be equal to the number of blocks.

  • radiuses (Sequence[Union[float, Sequence[float]]]) –

    Query radius for neighborhood grouping in each stage. The number of radiuses should be equal to the number of blocks + 1.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Maximum number of neighbors for each encoder stage. The number of num_neighbors should be equal to the number of blocks + 1.

  • add_self_loops (bool, default: False ) –

    Whether to include the center point as its own neighbor in grouping operations.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of point clouds (typically 3).

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the activation function.

  • act_first (bool, default: False ) –

    Whether to apply activation before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for normalization layer.

  • bias (Union[bool, List[bool]], default: True ) –

    Whether to use bias in linear / MLP layers.

  • dropout (float, default: 0.0 ) –

    Dropout probability before the classification head.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling operation for final feature aggregation.

Methods:

  • configure_stem –

    Build the stem MLP lifting the input features, or None when stem_channels is unset.

  • configure_encoder –

    Build the PointNeXtEncoder backbone.

  • configure_global_sa –

    Build the shared MLP over cat([pos, x]) applied before pooling, or None when global_sa_channels is unset.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) fed to the classification head, after the optional global set-abstraction.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) fed to the classification head, after the optional global set-abstraction.

configure_stem [source] ΒΆ

configure_stem() -> Optional[Module]

Build the stem MLP lifting the input features, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNeXtEncoder

Build the PointNeXtEncoder backbone.

configure_global_sa [source] ΒΆ

configure_global_sa() -> Optional[MLP]

Build the shared MLP over cat([pos, x]) applied before pooling, or None when global_sa_channels is unset.

PointNeXtPartSegmentation [source] ΒΆ

PointNeXtPartSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    num_categories: int,
    stem_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    stem_plain_last: bool = False,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_expansion: Union[int, Sequence[int]] = 4,
    sa_layers: int = 1,
    sa_use_res: bool = True,
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_plain_last: bool = True,
    ratios: Sequence[float],
    radiuses: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    add_self_loops: bool = False,
    spatial_dim: int = 3,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: Union[bool, List[bool]] = True,
    dropout: float = 0.0,
    head_channels: Optional[Sequence[int]] = None,
)

Bases: PartSegmentationModel

PointNeXt part segmentation model for ShapeNetPart.

Uses the same encoder as PointNeXtSegmentation but replaces the decoder with PointNeXtPartDecoder (global feature conditioning on shape category) and adds a head with global max+avg pooling.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of part classes (50 for ShapeNetPart).

  • num_categories (int) –

    Number of shape categories (16 for ShapeNetPart).

  • stem_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Stem MLP channel sizes.

  • encoder_channels (Sequence[int]) –

    Encoder channel dimensions per stage.

  • encoder_depths (Sequence[int]) –

    Residual block depths per encoder stage.

  • encoder_expansion (Union[int, Sequence[int]], default: 4 ) –

    InvResMLP expansion ratio.

  • sa_layers (int, default: 1 ) –

    Number of SA conv layers per block.

  • sa_use_res (bool, default: True ) –

    Whether SA blocks use residual connections.

  • decoder_channels (Sequence[int]) –

    Decoder channel dimensions per stage.

  • decoder_depths (Sequence[int]) –

    FP block depths per decoder stage.

  • ratios (Sequence[float]) –

    FPS downsampling ratios.

  • radiuses (Sequence[Union[float, Sequence[float]]]) –

    Ball-query radii.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Max neighbors per ball query.

Methods:

  • configure_stem –

    Build the stem MLP lifting the input features, or None when stem_channels is unset.

  • configure_encoder –

    Build the PointNeXtEncoder backbone.

  • configure_decoder –

    Build the PointNeXtPartDecoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: decoder features concatenated with their global max and mean pools.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: decoder features concatenated with their global max and mean pools.

configure_stem [source] ΒΆ

configure_stem() -> Optional[Module]

Build the stem MLP lifting the input features, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNeXtEncoder

Build the PointNeXtEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointNeXtPartDecoder

Build the PointNeXtPartDecoder upsampling the coarsest features back through the encoder skips.

PointNeXtSegmentation [source] ΒΆ

PointNeXtSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    stem_plain_last: bool = False,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_expansion: Union[int, Sequence[int]] = 4,
    sa_layers: int = 1,
    sa_use_res: bool = True,
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_plain_last: bool = True,
    ratios: Sequence[float],
    radiuses: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    add_self_loops: bool = False,
    spatial_dim: int = 3,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: Union[bool, List[bool]] = True,
    dropout: float = 0.0,
    head_channels: Optional[Sequence[int]] = None,
)

Bases: SemanticSegmentationModel

PointNeXt segmentation model as described in the paper PointNeXt: Revisiting PointNet++ with Improved Training and Scaling Strategies by Guocheng Qian, Yuchen Li, Houwen Peng, Jinjie Mai, Hasan Abed Al Kader Hammoud, Mohamed Elhoseiny, Bernard Ghanem.

PointNeXt modernizes PointNet++ through improved training strategies and architectural enhancements, achieving state-of-the-art performance while maintaining efficiency. The model introduces Inverted Residual MLP (InvResMLP) blocks, separable MLPs, and relative position normalization to enable effective model scaling.

Parameters:

  • in_channels (int) –

    Number of input channels (typically 3 for XYZ coordinates, or 6 for XYZ + RGB, or more with additional features like normal).

  • num_classes (int) –

    Number of output classes for segmentation.

  • stem_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Number of channels in the stem MLP layer(s) that map input to higher dimension. If None, no stem is used.

  • encoder_channels (Sequence[int]) –

    Channel dimensions for each encoder block. The number of channels should be equal to the number of blocks + 1.

  • encoder_depths (Sequence[int]) –

    Number of layers in each encoder block after the initial SA block.

  • encoder_expansion (Union[int, Sequence[int]], default: 4 ) –

    Expansion ratio to determine the hidden channels of the encoder blocks.

  • decoder_channels (Sequence[int]) –

    Channel dimensions for each decoder block. The number of channels should be equal to the number of decoder blocks + 1.

  • decoder_depths (Sequence[int]) –

    Number of layers in each decoder stage.

  • ratios (Sequence[float]) –

    Downsampling sampling ratios for each encoder stage. The number of ratios should be equal to the number of blocks.

  • radiuses (Sequence[Union[float, Sequence[float]]]) –

    Query radius for neighborhood grouping in each stage. The number of radiuses should be equal to the number of blocks + 1.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Maximum number of neighbors for each encoder stage. The number of num_neighbors should be equal to the number of blocks + 1.

  • add_self_loops (bool, default: False ) –

    Whether to include the center point as its own neighbor in grouping operations.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of point clouds (typically 3).

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for the activation function.

  • act_first (bool, default: False ) –

    Whether to apply activation before normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer type.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Additional arguments for normalization layer.

  • bias (Union[bool, List[bool]], default: True ) –

    Whether to use bias in linear / MLP layers.

  • dropout (float, default: 0.0 ) –

    Dropout probability before the classification head.

Methods:

  • configure_stem –

    Build the stem MLP lifting the input features, or None when stem_channels is unset.

  • configure_encoder –

    Build the PointNeXtEncoder backbone.

  • configure_decoder –

    Build the PointNeXtDecoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_stem [source] ΒΆ

configure_stem() -> Optional[Module]

Build the stem MLP lifting the input features, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNeXtEncoder

Build the PointNeXtEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PointNeXtDecoder

Build the PointNeXtDecoder upsampling the coarsest features back through the encoder skips.

PointPillarsDetection [source] ΒΆ

PointPillarsDetection(
    in_channels: int = 4,
    num_classes: int = 3,
    *,
    voxel_size: Sequence[float] = (0.16, 0.16, 4.0),
    point_cloud_range: Sequence[float] = (
        0.0,
        -39.68,
        -3.0,
        69.12,
        39.68,
        1.0,
    ),
    anchor_sizes: Sequence[Sequence[float]],
    anchor_bottom_heights: Sequence[float],
    feature_map_stride: int,
    anchor_rotations: Sequence[float] = (0.0, 1.57),
    feat_channels: Sequence[int] = (64,),
    layer_nums: Sequence[int] = (3, 5, 5),
    layer_strides: Sequence[int] = (2, 2, 2),
    num_filters: Sequence[int] = (64, 128, 256),
    upsample_strides: Sequence[int] = (1, 2, 4),
    num_upsample_filters: Sequence[int] = (128, 128, 128),
    num_dir_bins: int = 2,
    dir_offset: float = 0.78539,
    dir_limit_offset: 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: DetectionModel

PointPillars 3D object detector (packed point format).

Reference: PointPillars: Fast Encoders for Object Detection from Point Clouds (Lang et al., 2019). Reference implementation: open-mmlab/OpenPCDet.

Parameters:

  • in_channels (int, default: 4 ) –

    Raw point feature channels including xyz (e.g. \(4\) for \(x, y, z, \text{intensity}\)).

  • num_classes (int, default: 3 ) –

    Number of foreground classes.

  • voxel_size (Sequence[float], default: (0.16, 0.16, 4.0) ) –

    Pillar size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (0.0, -39.68, -3.0, 69.12, 39.68, 1.0) ) –

    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)\), one row per class.

  • anchor_bottom_heights (Sequence[float]) –

    Per-class anchor bottom \(z\), one per class.

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

  • feat_channels (Sequence[int], default: (64,) ) –

    Output channels of the pillar feature net.

  • layer_nums (Sequence[int], default: (3, 5, 5) ) –

    2D backbone conv counts per level.

  • layer_strides (Sequence[int], default: (2, 2, 2) ) –

    2D backbone downsample strides per level.

  • num_filters (Sequence[int], default: (64, 128, 256) ) –

    2D backbone channel widths per level.

  • upsample_strides (Sequence[int], default: (1, 2, 4) ) –

    2D backbone upsample strides per level.

  • num_upsample_filters (Sequence[int], default: (128, 128, 128) ) –

    2D backbone upsample channels per level.

  • num_dir_bins (int, default: 2 ) –

    Number of direction bins in the head.

  • dir_offset (float, default: 0.78539 ) –

    Direction-classifier angle offset.

  • dir_limit_offset (float, default: 0.0 ) –

    Heading wrap offset used during decoding.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation type or callable for the pillar feature net and 2D backbone.

  • 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 pillar feature net and 2D backbone.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra normalization arguments.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_vfe –

    Build the pillar feature net.

  • configure_backbone –

    Build the 2D BEV backbone.

  • configure_head –

    Build the single-group anchor head.

  • decode –

    Decode a forward output into raw per-anchor detections (see AnchorHead.decode).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the BEV feature map entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_vfe [source] ΒΆ

configure_vfe() -> PillarFeatureNet

Build the pillar feature net.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVBackbone

Build the 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> AnchorHead

Build the single-group anchor head.

decode [source] ΒΆ

decode(out: AnchorHeadOutput) -> Detection3D

Decode a forward output into raw per-anchor detections (see AnchorHead.decode).

PointPillarsMultiHeadDetection [source] ΒΆ

PointPillarsMultiHeadDetection(
    in_channels: int = 5,
    num_classes: int = 10,
    *,
    voxel_size: Sequence[float] = (0.2, 0.2, 8.0),
    point_cloud_range: Sequence[float] = (
        -51.2,
        -51.2,
        -5.0,
        51.2,
        51.2,
        3.0,
    ),
    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),
    feat_channels: Sequence[int] = (64,),
    layer_nums: Sequence[int] = (3, 5, 5),
    layer_strides: Sequence[int] = (2, 2, 2),
    num_filters: Sequence[int] = (64, 128, 256),
    upsample_strides: Sequence[float] = (0.5, 1, 2),
    num_upsample_filters: Sequence[int] = (128, 128, 128),
    shared_conv_num_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: DetectionModel

PointPillars with a multi-group anchor head (nuScenes 10-class, packed point format).

Reference implementation: open-mmlab/OpenPCDet (cbgs_pp_multihead). Same pillar trunk as PointPillarsDetection but with an MultiGroupAnchorHead head (per-group heads, sincos + velocity box code). Input points carry 5 features (\(x, y, z, \text{intensity}, \Delta t\) from 10-sweep aggregation).

Parameters:

  • in_channels (int, default: 5 ) –

    Raw point feature channels including xyz (5 for nuScenes).

  • num_classes (int, default: 10 ) –

    Number of foreground classes (10 for nuScenes).

  • voxel_size (Sequence[float], default: (0.2, 0.2, 8.0) ) –

    Pillar size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (-51.2, -51.2, -5.0, 51.2, 51.2, 3.0) ) –

    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)\), one row per class.

  • anchor_bottom_heights (Sequence[float]) –

    Per-class anchor bottom \(z\), one per class.

  • head_class_groups (Sequence[Sequence[int]]) –

    Class-index groups, one per RPN head (e.g. [[0], [1, 2], ...]).

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

  • feat_channels (Sequence[int], default: (64,) ) –

    Pillar feature-net output channels.

  • layer_nums (Sequence[int], default: (3, 5, 5) ) –

    2D backbone conv counts per level.

  • layer_strides (Sequence[int], default: (2, 2, 2) ) –

    2D backbone downsample strides per level.

  • num_filters (Sequence[int], default: (64, 128, 256) ) –

    2D backbone channel widths per level.

  • upsample_strides (Sequence[float], default: (0.5, 1, 2) ) –

    2D backbone upsample factors per level (may be < 1).

  • num_upsample_filters (Sequence[int], default: (128, 128, 128) ) –

    2D backbone upsample channels per level.

  • shared_conv_num_filter (int, default: 64 ) –

    Channels of the head's shared conv.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation type or callable for the pillar feature net, 2D backbone and head.

  • 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 pillar feature net, 2D backbone and head.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra normalization arguments.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_vfe –

    Build the pillar feature net.

  • configure_backbone –

    Build the 2D BEV backbone.

  • configure_head –

    Build the multi-group anchor head.

  • decode –

    Decode a forward output into raw per-anchor detections (see MultiGroupAnchorHead.decode).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the BEV feature map entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_vfe [source] ΒΆ

configure_vfe() -> PillarFeatureNet

Build the pillar feature net.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVBackbone

Build the 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> MultiGroupAnchorHead

Build the multi-group anchor head.

decode [source] ΒΆ

Decode a forward output into raw per-anchor detections (see MultiGroupAnchorHead.decode).

PointRCNNDetection [source] ΒΆ

PointRCNNDetection(
    in_channels: int = 4,
    num_classes: int = 3,
    *,
    mean_sizes: Union[Tensor, Sequence[Sequence[float]]],
    sa_channels: Sequence[Sequence[Sequence[int]]],
    sa_num_points: Sequence[int],
    sa_radii: Sequence[Sequence[float]],
    sa_num_neighbors: Sequence[Sequence[int]],
    fp_channels: Sequence[Sequence[int]],
    point_cls_channels: Sequence[int] = (256, 256),
    point_reg_channels: Sequence[int] = (256, 256),
    roi_sa_channels: Sequence[Sequence[int]],
    roi_sa_num_points: Sequence[int],
    roi_sa_radii: Sequence[float],
    roi_sa_num_neighbors: Sequence[int],
    roi_xyz_up_channels: Sequence[int] = (128, 128),
    roi_cls_channels: Sequence[int] = (256, 256),
    roi_reg_channels: Sequence[int] = (256, 256),
    num_sampled_points: int = 512,
    pool_extra_width: Sequence[float] = (0.0, 0.0, 0.0),
    depth_normalizer: float = 70.0,
    proposal_pre_maxsize: int = 9000,
    proposal_post_maxsize: int = 100,
    proposal_nms_thresh: float = 0.85,
    proposal_nms_rotated: bool = True,
    train_proposal_post_maxsize: int = 512,
    train_proposal_nms_thresh: float = 0.8,
    roi_per_image: int = 128,
    fg_ratio: float = 0.5,
    reg_fg_thresh: float = 0.55,
    cls_fg_thresh: float = 0.6,
    cls_bg_thresh_lo: float = 0.1,
    hard_bg_ratio: float = 0.8,
    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: DetectionModel

PointRCNN two-stage point-based 3D object detector (packed point format).

Reference: PointRCNN: 3D Object Proposal Generation and Detection from Point Cloud (Shi et al., 2019). Reference implementation: open-mmlab/OpenPCDet.

Stage 1 runs a multi-scale PointNet++ U-Net (PointNet2Encoder + PointNet2Decoder) over the raw point cloud, then a per-point head (PointHeadBox) predicts foreground scores and one box proposal per point. The top proposals (after class-agnostic NMS) become ROIs that stage 2 (PointRCNNRefinementHead) pools points around, canonically transforms, and refines into a confidence and a box correction.

Parameters:

  • in_channels (int, default: 4 ) –

    Raw point feature channels including xyz (e.g. \(4\) for \(x, y, z, \text{intensity}\)).

  • num_classes (int, default: 3 ) –

    Number of foreground classes.

  • mean_sizes (Union[Tensor, Sequence[Sequence[float]]]) –

    Per-class mean box size \((d_x, d_y, d_z)\), shape \((\text{num\_classes}, 3)\).

  • sa_channels (Sequence[Sequence[Sequence[int]]]) –

    Stage-1 per-SA-block, per-scale MLP channel lists.

  • sa_num_points (Sequence[int]) –

    Stage-1 per-SA-block sample counts.

  • sa_radii (Sequence[Sequence[float]]) –

    Stage-1 per-SA-block, per-scale ball-query radii.

  • sa_num_neighbors (Sequence[Sequence[int]]) –

    Stage-1 per-SA-block, per-scale neighbor caps.

  • fp_channels (Sequence[Sequence[int]]) –

    Stage-1 per-FP-block MLP channel lists, ordered from the coarsest skip level to the finest (PointNet2Decoder order).

  • point_cls_channels (Sequence[int], default: (256, 256) ) –

    Stage-1 classification MLP hidden channels.

  • point_reg_channels (Sequence[int], default: (256, 256) ) –

    Stage-1 box-regression MLP hidden channels.

  • roi_sa_channels (Sequence[Sequence[int]]) –

    Stage-2 per-SA-block MLP channel lists.

  • roi_sa_num_points (Sequence[int]) –

    Stage-2 per-SA-block sample counts (-1 groups all).

  • roi_sa_radii (Sequence[float]) –

    Stage-2 per-SA-block ball-query radii.

  • roi_sa_num_neighbors (Sequence[int]) –

    Stage-2 per-SA-block neighbor caps.

  • roi_xyz_up_channels (Sequence[int], default: (128, 128) ) –

    Stage-2 canonical-xyz lifting MLP channels.

  • roi_cls_channels (Sequence[int], default: (256, 256) ) –

    Stage-2 confidence MLP hidden channels.

  • roi_reg_channels (Sequence[int], default: (256, 256) ) –

    Stage-2 box-refinement MLP hidden channels.

  • num_sampled_points (int, default: 512 ) –

    Points pooled per ROI in stage 2.

  • pool_extra_width (Sequence[float], default: (0.0, 0.0, 0.0) ) –

    Per-axis enlargement of the stage-2 pooling box.

  • depth_normalizer (float, default: 70.0 ) –

    Divisor for the stage-2 point-depth feature.

  • proposal_pre_maxsize (int, default: 9000 ) –

    Proposals kept before stage-1 NMS.

  • proposal_post_maxsize (int, default: 100 ) –

    ROIs kept after stage-1 NMS at inference (the stage-2 batch size per scene).

  • proposal_nms_thresh (float, default: 0.85 ) –

    Stage-1 proposal NMS IoU threshold at inference.

  • proposal_nms_rotated (bool, default: True ) –

    Run the stage-1 NMS on the rotated BEV IoU (the reference nms_gpu) rather than the axis-aligned 3D IoU.

  • train_proposal_post_maxsize (int, default: 512 ) –

    Proposals kept after stage-1 NMS during training (before ROI sampling).

  • train_proposal_nms_thresh (float, default: 0.8 ) –

    Stage-1 proposal NMS IoU threshold during training.

  • roi_per_image (int, default: 128 ) –

    ROIs sampled per scene for stage-2 training.

  • fg_ratio (float, default: 0.5 ) –

    Target fraction of foreground ROIs in the sampled set.

  • reg_fg_thresh (float, default: 0.55 ) –

    ROI-to-GT IoU at or above which a sampled ROI is foreground (box regression valid).

  • cls_fg_thresh (float, default: 0.6 ) –

    ROI-to-GT IoU used with reg_fg_thresh to define the foreground sampling threshold.

  • cls_bg_thresh_lo (float, default: 0.1 ) –

    ROI-to-GT IoU below which a ROI is easy background (else hard background).

  • hard_bg_ratio (float, default: 0.8 ) –

    Fraction of the sampled background ROIs drawn from hard (higher-IoU) background.

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

Methods:

  • configure_encoder –

    Build the stage-1 multi-scale PointNet++ encoder.

  • configure_decoder –

    Build the stage-1 PointNet++ feature-propagation decoder.

  • configure_point_head –

    Build the stage-1 per-point foreground and box-proposal head.

  • configure_roi_head –

    Build the stage-2 ROI refinement head.

  • forward –

    Run both stages; in train mode the ground truth drives the stage-2 ROI sampling.

  • decode –

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

Attributes:

  • num_features (int) –

    Channel count \(C\) of the per-point backbone features entering the heads.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the per-point backbone features entering the heads.

configure_encoder [source] ΒΆ

configure_encoder() -> PointNet2Encoder

Build the stage-1 multi-scale PointNet++ encoder.

configure_decoder [source] ΒΆ

configure_decoder() -> PointNet2Decoder

Build the stage-1 PointNet++ feature-propagation decoder.

configure_point_head [source] ΒΆ

configure_point_head() -> PointHeadBox

Build the stage-1 per-point foreground and box-proposal head.

configure_roi_head [source] ΒΆ

configure_roi_head() -> PointRCNNRefinementHead

Build the stage-2 ROI refinement head.

forward [source] ΒΆ

forward(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    gt_boxes: OptTensor = None,
    gt_labels: OptTensor = None,
    gt_batch: OptTensor = None,
) -> Union[PointRCNNTrainOutput, PointRCNNOutput]

Run both stages; in train mode the ground truth drives the stage-2 ROI sampling.

Stage-2 training samples its ROIs by matching stage-1 proposals to ground-truth boxes at forward time, so train mode requires the packed ground truth. The GT arguments default to None and are omitted at inference (a training pipeline passes box / label / batch_box after the point inputs, e.g. via input_keys).

Parameters:

  • x (OptTensor) –

    Per-point features including reflectance, shape \((N, \text{in\_channels} - 3)\).

  • pos (Tensor) –

    Per-point coordinates, shape \((N, 3)\).

  • batch (Tensor) –

    Per-point scene index, shape \((N,)\).

  • gt_boxes (OptTensor, default: None ) –

    Ground-truth boxes \((K, 7)\), required in train mode.

  • gt_labels (OptTensor, default: None ) –

    Ground-truth \(0\)-based classes, shape \((K,)\), required in train mode.

  • gt_batch (OptTensor, default: None ) –

    Per-box scene index, shape \((K,)\), required in train mode.

Returns:

Shape
  • x: \((N, \text{in\_channels} - 3)\), pos: \((N, 3)\), batch: \((N,)\)
  • gt_boxes: \((K, 7)\), gt_labels / gt_batch: \((K,)\)

decode [source] ΒΆ

decode(out: PointRCNNOutput) -> Detection3D

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

Scores each refined box by its stage-2 confidence (sigmoid of rcnn_cls) and labels it by the stage-1 ROI label (shifted to 0-indexed). The full per-ROI set is returned; the evaluation pipeline applies class-agnostic 3D NMS then score thresholding via the torch_pointcloud.ops.box3d utilities (see the benchmark example).

Parameters:

  • out (PointRCNNOutput) –

    A forward output {"rcnn_cls", "boxes", "roi_labels", "roi_scores", "batch"}.

Returns:

  • Detection3D –

    Packed per-ROI detections {"boxes": (R, 7), "scores": (R,), "labels": (R,), "batch": (R,)}

  • Detection3D –

    (PyG layout).

PVCNNClassification [source] ΒΆ

PVCNNClassification(
    in_channels: int,
    num_classes: int,
    *,
    channels: Sequence[int],
    global_channels: Optional[Sequence[int]] = None,
    depths: Sequence[int],
    kernel_sizes: Sequence[int],
    resolutions: Sequence[int],
    use_se: bool = False,
    normalize: bool = True,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    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,
)

Bases: ClassificationModel

PVCNN classification model from Point-Voxel CNN for Efficient 3D Deep Learning by Zhijian Liu, Haotian Tang, Yujun Lin, Song Han.

A flat stack of point-voxel conv blocks at decreasing voxel resolutions, followed by global pooling, an optional global MLP and a linear head.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of output classes. \(0\) replaces the head with nn.Identity.

  • channels (Sequence[int]) –

    Output feature width of each block.

  • global_channels (Optional[Sequence[int]], default: None ) –

    MLP channels applied to the pooled feature. None skips the global MLP.

  • depths (Sequence[int]) –

    Number of point-voxel conv layers per block.

  • kernel_sizes (Sequence[int]) –

    Voxel conv kernel size per block.

  • resolutions (Sequence[int]) –

    Voxel grid resolution per block; \(0\) or None uses MLP layers instead.

  • use_se (bool, default: False ) –

    Add a squeeze-and-excitation gate to every voxel branch.

  • normalize (bool, default: True ) –

    Normalize coordinates into the unit voxel grid before voxelization.

  • dropout (float, default: 0.0 ) –

    Dropout applied to the pooled features before the head.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling reducing the point features to one vector per cloud.

  • 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 ) –

    Apply the activation before the normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization function.

Methods:

  • configure_blocks –

    Builds the point-voxel conv blocks, one per entry of depths.

  • configure_global_mlp –

    Builds the MLP applied to the pooled global feature, or None when global_channels is empty.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the pooled features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the pooled features entering the head.

configure_blocks [source] ΒΆ

configure_blocks() -> ModuleList

Builds the point-voxel conv blocks, one per entry of depths.

configure_global_mlp [source] ΒΆ

configure_global_mlp() -> Optional[MLP]

Builds the MLP applied to the pooled global feature, or None when global_channels is empty.

PVCNNSegmentation [source] ΒΆ

PVCNNSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    channels: Sequence[int],
    global_channels: Optional[Sequence[int]] = None,
    depths: Sequence[int],
    kernel_sizes: Sequence[int],
    resolutions: Sequence[int],
    spatial_dim: int = 3,
    use_se: bool = False,
    normalize: bool = True,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    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,
    head_channels: Optional[Sequence[int]] = None,
    head_dropout: float = 0.0,
)

Bases: SemanticSegmentationModel

PVCNN segmentation model from Point-Voxel CNN for Efficient 3D Deep Learning by Zhijian Liu, Haotian Tang, Yujun Lin, Song Han.

The same flat stack of point-voxel conv blocks as PVCNNClassification, but every layer output is kept at full point resolution and concatenated with the broadcast global feature, so no decoder is needed to recover per-point logits.

Parameters:

  • in_channels (int) –

    Number of input feature channels. Raised to spatial_dim when smaller.

  • num_classes (int) –

    Number of output classes. \(0\) replaces the head with nn.Identity.

  • channels (Sequence[int]) –

    Output feature width of each block.

  • global_channels (Optional[Sequence[int]], default: None ) –

    MLP channels applied to the pooled feature. None skips the global MLP.

  • depths (Sequence[int]) –

    Number of point-voxel conv layers per block.

  • kernel_sizes (Sequence[int]) –

    Voxel conv kernel size per block.

  • resolutions (Sequence[int]) –

    Voxel grid resolution per block; \(0\) or None uses MLP layers instead.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of the point clouds.

  • use_se (bool, default: False ) –

    Add a squeeze-and-excitation gate to every voxel branch.

  • normalize (bool, default: True ) –

    Normalize coordinates into the unit voxel grid before voxelization.

  • dropout (float, default: 0.0 ) –

    Dropout applied to the concatenated point features before the head.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling producing the cloud-level feature that is broadcast back to the points.

  • 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 ) –

    Apply the activation before the normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization function.

  • head_channels (Optional[Sequence[int]], default: None ) –

    Hidden MLP channels of the head. None uses a single linear layer.

  • head_dropout (float, default: 0.0 ) –

    Dropout between the hidden head layers.

Methods:

  • configure_blocks –

    Builds the point-voxel conv blocks, one per entry of depths.

  • configure_global_mlp –

    Builds the MLP applied to the pooled global feature, or None when global_channels is empty.

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: every layer output concatenated with the global feature.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: every layer output concatenated with the global feature.

configure_blocks [source] ΒΆ

configure_blocks() -> ModuleList

Builds the point-voxel conv blocks, one per entry of depths.

configure_global_mlp [source] ΒΆ

configure_global_mlp() -> Optional[MLP]

Builds the MLP applied to the pooled global feature, or None when global_channels is empty.

PVCNN2Classification [source] ΒΆ

PVCNN2Classification(
    in_channels: int,
    num_classes: int,
    *,
    ratios: Sequence[float],
    radii: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    sa_channels: Sequence[
        Sequence[Union[int, Sequence[int]]]
    ],
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_resolutions: Sequence[Optional[int]],
    encoder_kernel_sizes: Sequence[int],
    use_se: bool = False,
    normalize: bool = True,
    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,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
)

Bases: ClassificationModel

PVCNN++ classification model from Point-Voxel CNN for Efficient 3D Deep Learning by Zhijian Liu, Haotian Tang, Yujun Lin, Song Han.

PVCNN++ grafts point-voxel convolutions onto a PointNet++-style hierarchy: every encoder block downsamples the cloud with a set-abstraction module (farthest point sampling + ball grouping), then refines the abstracted features with PVConv layers that fuse a coarse voxel-grid convolution branch with a per-point MLP branch. The final features are pooled into a global embedding and classified with a linear head.

Parameters:

  • in_channels (int) –

    Number of input feature channels. Takes precedence over encoder_channels[0] when the two disagree.

  • num_classes (int) –

    Number of output classes. \(0\) replaces the head with nn.Identity.

  • ratios (Sequence[float]) –

    Farthest-point-sampling ratio per encoder block. Block \(i\) uses ratios[i - 1] (block \(0\) has no set-abstraction module), so the list pads to the block count and its final entry is unused.

  • radii (Sequence[Union[float, Sequence[float]]]) –

    Ball-query radius per encoder block, aligned like ratios. A nested sequence enables multi-scale grouping.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Maximum number of neighbors per ball query, aligned like ratios.

  • sa_channels (Sequence[Sequence[Union[int, Sequence[int]]]]) –

    MLP channels of each set-abstraction module, aligned like ratios. Nest twice for multi-scale grouping.

  • encoder_channels (Sequence[int]) –

    Per-block feature widths, one more entry than the number of blocks: encoder_channels[i + 1] is the output width of block \(i\).

  • encoder_depths (Sequence[int]) –

    Number of point-voxel conv layers per block; \(0\) keeps the block set-abstraction-only.

  • encoder_resolutions (Sequence[Optional[int]]) –

    Voxel grid resolution per block; \(0\) or None uses MLP layers instead of point-voxel convs.

  • encoder_kernel_sizes (Sequence[int]) –

    Voxel conv kernel size per block.

  • use_se (bool, default: False ) –

    Add a squeeze-and-excitation gate to every voxel branch.

  • normalize (bool, default: True ) –

    Normalize coordinates into the unit voxel grid before voxelization.

  • 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 ) –

    Apply the activation before the normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization function.

  • dropout (float, default: 0.0 ) –

    Dropout probability applied to the pooled global embedding.

  • global_pool (PoolLike, default: 'max' ) –

    Pooling used to aggregate per-point features into the global embedding.

Shape
  • x: \((N, C_{in})\) point features; when None, pos is used as features.
  • pos: \((N, 3)\) point coordinates.
  • batch: \((N,)\) batch indices.
  • output: \((B, \text{num\_classes})\) classification logits.

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PVCNN2Encoder

Build the PVCNN2Encoder backbone.

PVCNN2Segmentation [source] ΒΆ

PVCNN2Segmentation(
    in_channels: int,
    num_classes: int,
    *,
    ratios: Sequence[float],
    radii: Sequence[Union[float, Sequence[float]]],
    num_neighbors: Sequence[Union[int, Sequence[int]]],
    sa_channels: Sequence[
        Sequence[Union[int, Sequence[int]]]
    ],
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_resolutions: Sequence[Optional[int]],
    encoder_kernel_sizes: Sequence[int],
    fp_channels: Sequence[Sequence[int]],
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_resolutions: Sequence[Optional[int]],
    decoder_kernel_sizes: Sequence[int],
    use_se: bool = False,
    normalize: bool = True,
    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,
    dropout: float = 0.0,
    head_channels: Optional[Sequence[int]] = None,
    head_dropout: float = 0.0,
)

Bases: SemanticSegmentationModel

PVCNN++ segmentation model from Point-Voxel CNN for Efficient 3D Deep Learning by Zhijian Liu, Haotian Tang, Yujun Lin, Song Han.

A U-shaped PVCNN++: the encoder alternates set-abstraction downsampling with point-voxel conv refinement, and the decoder upsamples back with feature-propagation blocks followed by point-voxel convs. Decoder skips are consumed deepest-first, and the last decoder block uses the extra input features (the input without its 3 leading coordinate channels) as its skip, so the encoder may have more blocks than the decoder (the leftover shallow intermediates are skipped).

Parameters:

  • in_channels (int) –

    Number of input feature channels. Takes precedence over encoder_channels[0] when the two disagree; \(0\) falls back to \(3\) (positions used as features).

  • num_classes (int) –

    Number of output classes. \(0\) replaces the head with nn.Identity.

  • ratios (Sequence[float]) –

    Farthest-point-sampling ratio per encoder block. Block \(i\) uses ratios[i - 1] (block \(0\) has no set-abstraction module), so the list pads to the block count and its final entry is unused.

  • radii (Sequence[Union[float, Sequence[float]]]) –

    Ball-query radius per encoder block, aligned like ratios. A nested sequence enables multi-scale grouping.

  • num_neighbors (Sequence[Union[int, Sequence[int]]]) –

    Maximum number of neighbors per ball query, aligned like ratios.

  • sa_channels (Sequence[Sequence[Union[int, Sequence[int]]]]) –

    MLP channels of each set-abstraction module, aligned like ratios. Nest twice for multi-scale grouping.

  • encoder_channels (Sequence[int]) –

    Per-block feature widths, one more entry than the number of encoder blocks: encoder_channels[i + 1] is the output width of block \(i\).

  • encoder_depths (Sequence[int]) –

    Number of point-voxel conv layers per encoder block; \(0\) keeps the block set-abstraction-only.

  • encoder_resolutions (Sequence[Optional[int]]) –

    Voxel grid resolution per encoder block; \(0\) or None uses MLP layers instead of point-voxel convs.

  • encoder_kernel_sizes (Sequence[int]) –

    Voxel conv kernel size per encoder block.

  • fp_channels (Sequence[Sequence[int]]) –

    MLP channels of each feature-propagation module, deepest stage first.

  • decoder_channels (Sequence[int]) –

    Per-block feature widths of the decoder conv stacks.

  • decoder_depths (Sequence[int]) –

    Number of point-voxel conv layers per decoder block; \(0\) keeps the block feature-propagation-only. At most one decoder block per set-abstraction level (encoder blocks minus one).

  • decoder_resolutions (Sequence[Optional[int]]) –

    Voxel grid resolution per decoder block; \(0\) or None uses MLP layers instead of point-voxel convs.

  • decoder_kernel_sizes (Sequence[int]) –

    Voxel conv kernel size per decoder block.

  • use_se (bool, default: False ) –

    Add a squeeze-and-excitation gate to every voxel branch.

  • normalize (bool, default: True ) –

    Normalize coordinates into the unit voxel grid before voxelization.

  • 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 ) –

    Apply the activation before the normalization.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization function.

  • dropout (float, default: 0.0 ) –

    Dropout probability before the linear head. Ignored when head_channels is set: the MLP head applies its own head_dropout.

  • head_channels (Optional[Sequence[int]], default: None ) –

    Hidden widths of the segmentation head MLP. None uses a single linear layer.

  • head_dropout (float, default: 0.0 ) –

    Dropout probability inside the head MLP.

Shape
  • x: \((N, C_{in})\) point features; when None, pos is used as features.
  • pos: \((N, 3)\) point coordinates.
  • batch: \((N,)\) batch indices.
  • output: \((N, \text{num\_classes})\) per-point segmentation logits.

Methods:

  • configure_encoder –

    Build the PVCNN2Encoder backbone.

  • configure_decoder –

    Build the PVCNN2Decoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_encoder [source] ΒΆ

configure_encoder() -> PVCNN2Encoder

Build the PVCNN2Encoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> PVCNN2Decoder

Build the PVCNN2Decoder upsampling the coarsest features back through the encoder skips.

RandLANetClassification [source] ΒΆ

RandLANetClassification(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Optional[int] = 8,
    encoder_channels: Sequence[int],
    decimation: Union[int, Sequence[int]] = 4,
    num_neighbors: Union[int, Sequence[int]] = 16,
    aggr_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    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,
)

Bases: ClassificationModel

RandLA-Net classification model from RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds by Qingyong Hu, Bo Yang, Linhai Xie, Stefano Rosa, Yulan Guo, Zhihua Wang, Niki Trigoni, Andrew Markham.

Random sampling for downsampling and dilated residual blocks with local feature aggregation; point features are pooled globally after the encoder for classification.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of classes.

  • stem_channels (Optional[int], default: 8 ) –

    Number of channels in the stem MLP. Set to None to skip the stem.

  • encoder_channels (Sequence[int]) –

    Output channels of each dilated residual block (must be even).

  • decimation (Union[int, Sequence[int]], default: 4 ) –

    Decimation factor between consecutive encoder blocks.

  • num_neighbors (Union[int, Sequence[int]], default: 16 ) –

    Number of neighbors for the kNN graph in each block.

  • aggr_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Optional channels for the aggregation MLP applied before global pooling.

  • dropout (float, default: 0.0 ) –

    Dropout rate before the classification head.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling operation.

Methods:

  • configure_stem –

    Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

  • configure_encoder –

    Build the RandLANetEncoder backbone.

  • configure_aggr –

    Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output, after the optional aggregation MLP.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output, after the optional aggregation MLP.

configure_stem [source] ΒΆ

configure_stem() -> Optional[MLP]

Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> RandLANetEncoder

Build the RandLANetEncoder backbone.

configure_aggr [source] ΒΆ

configure_aggr() -> Optional[MLP]

Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

RandLANetSegmentation [source] ΒΆ

RandLANetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Optional[int] = 8,
    encoder_channels: Sequence[int],
    fp_channels: Sequence[int],
    head_channels: Sequence[int] = (64, 32),
    decimation: Union[int, Sequence[int]] = 4,
    num_neighbors: Union[int, Sequence[int]] = 16,
    aggr_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    dropout: float = 0.5,
    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,
)

Bases: SemanticSegmentationModel

RandLA-Net segmentation model from RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds by Qingyong Hu, Bo Yang, Linhai Xie, Stefano Rosa, Yulan Guo, Zhihua Wang, Niki Trigoni, Andrew Markham.

Encoder uses random sampling between dilated residual blocks; decoder uses 1-NN nearest-neighbor interpolation with concatenation skips and a single 1x1 linear+BN+act per stage. The skip used at full resolution is the pre-decimation output of the first encoder block.

Parameters:

  • in_channels (int) –

    Number of input channels.

  • num_classes (int) –

    Number of classes.

  • stem_channels (Optional[int], default: 8 ) –

    Number of channels in the stem MLP. Set to None to skip the stem.

  • encoder_channels (Sequence[int]) –

    Output channels of each dilated residual block (must be even).

  • fp_channels (Sequence[int]) –

    Per-stage decoder channels (one list per upsampling step).

  • head_channels (Sequence[int], default: (64, 32) ) –

    Hidden channels of the segmentation head MLP.

  • decimation (Union[int, Sequence[int]], default: 4 ) –

    Decimation factor between consecutive encoder blocks.

  • num_neighbors (Union[int, Sequence[int]], default: 16 ) –

    Number of neighbors for the kNN graph in each block.

  • aggr_channels (Optional[Union[int, Sequence[int]]], default: None ) –

    Channels for the bottleneck MLP between the encoder and the decoder (the upstream "decoder_0").

  • dropout (float, default: 0.5 ) –

    Dropout rate inside the segmentation head MLP.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation type for both the decoder FP MLPs and the segmentation head MLP (string passed to create_act, or a Callable / nn.Module).

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments forwarded to the activation.

  • act_first (bool, default: False ) –

    If True, apply activation before normalization (PyG's MLP act_first semantics): Linear β†’ Act β†’ Norm β†’ Dropout instead of the default Linear β†’ Norm β†’ Act β†’ Dropout.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization type for the decoder FP MLPs and the head MLP.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments forwarded to the normalization layers.

  • bias (bool, default: False ) –

    Whether the decoder/head hidden linear layers carry an explicit bias. The final head layer always uses bias=True since it has no normalization.

Methods:

  • configure_stem –

    Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

  • configure_encoder –

    Build the RandLANetEncoder backbone.

  • configure_aggr –

    Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

  • configure_decoder –

    Build the RandLANetDecoder upsampling the bottleneck features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_stem [source] ΒΆ

configure_stem() -> Optional[MLP]

Build the stem lifting the input features to stem_channels, or None when stem_channels is unset.

configure_encoder [source] ΒΆ

configure_encoder() -> RandLANetEncoder

Build the RandLANetEncoder backbone.

configure_aggr [source] ΒΆ

configure_aggr() -> Optional[MLP]

Build the aggregation MLP applied to the encoder output, or None when aggr_channels is unset.

configure_decoder [source] ΒΆ

configure_decoder() -> RandLANetDecoder

Build the RandLANetDecoder upsampling the bottleneck features back through the encoder skips.

SECONDDetection [source] ΒΆ

SECONDDetection(
    in_channels: int = 4,
    num_classes: int = 3,
    *,
    voxel_size: Sequence[float] = (0.05, 0.05, 0.1),
    point_cloud_range: Sequence[float] = (
        0.0,
        -40.0,
        -3.0,
        70.4,
        40.0,
        1.0,
    ),
    anchor_sizes: Sequence[Sequence[float]],
    anchor_bottom_heights: Sequence[float],
    feature_map_stride: int,
    anchor_rotations: Sequence[float] = (0.0, 1.57),
    layer_nums: Sequence[int] = (5, 5),
    layer_strides: Sequence[int] = (1, 2),
    num_filters: Sequence[int] = (128, 256),
    upsample_strides: Sequence[int] = (1, 2),
    num_upsample_filters: Sequence[int] = (256, 256),
    num_dir_bins: int = 2,
    dir_offset: float = 0.78539,
    dir_limit_offset: 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: DetectionModel

SECOND 3D object detector (packed point format).

Reference: SECOND: Sparsely Embedded Convolutional Detection (Yan et al., 2018). Reference implementation: open-mmlab/OpenPCDet.

Parameters:

  • in_channels (int, default: 4 ) –

    Raw point feature channels including xyz (e.g. \(4\) for \(x, y, z, \text{intensity}\)).

  • num_classes (int, default: 3 ) –

    Number of foreground classes.

  • voxel_size (Sequence[float], default: (0.05, 0.05, 0.1) ) –

    Voxel size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (0.0, -40.0, -3.0, 70.4, 40.0, 1.0) ) –

    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)\), one row per class.

  • anchor_bottom_heights (Sequence[float]) –

    Per-class anchor bottom \(z\), one per class.

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

  • layer_nums (Sequence[int], default: (5, 5) ) –

    2D backbone conv counts per level.

  • layer_strides (Sequence[int], default: (1, 2) ) –

    2D backbone downsample strides per level.

  • num_filters (Sequence[int], default: (128, 256) ) –

    2D backbone channel widths per level.

  • upsample_strides (Sequence[int], default: (1, 2) ) –

    2D backbone upsample strides per level.

  • num_upsample_filters (Sequence[int], default: (256, 256) ) –

    2D backbone upsample channels per level.

  • num_dir_bins (int, default: 2 ) –

    Number of direction bins in the head.

  • dir_offset (float, default: 0.78539 ) –

    Direction-classifier angle offset.

  • dir_limit_offset (float, default: 0.0 ) –

    Heading wrap offset used during decoding.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation type or callable for the 3D/2D backbones.

  • 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 3D/2D backbones.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra normalization arguments.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_backbone_3d –

    Build the sparse 3D voxel backbone.

  • configure_backbone –

    Build the 2D BEV backbone.

  • configure_head –

    Build the single-group anchor head.

  • decode –

    Decode a forward output into raw per-anchor detections (see AnchorHead.decode).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the BEV feature map entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_backbone_3d [source] ΒΆ

configure_backbone_3d() -> VoxelBackbone8x

Build the sparse 3D voxel backbone.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVBackbone

Build the 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> AnchorHead

Build the single-group anchor head.

decode [source] ΒΆ

decode(out: AnchorHeadOutput) -> Detection3D

Decode a forward output into raw per-anchor detections (see AnchorHead.decode).

SECONDMultiHeadDetection [source] ΒΆ

SECONDMultiHeadDetection(
    in_channels: int = 5,
    num_classes: int = 10,
    *,
    voxel_size: Sequence[float] = (0.1, 0.1, 0.2),
    point_cloud_range: Sequence[float] = (
        -51.2,
        -51.2,
        -5.0,
        51.2,
        51.2,
        3.0,
    ),
    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),
    layer_nums: Sequence[int] = (5, 5),
    layer_strides: Sequence[int] = (1, 2),
    num_filters: Sequence[int] = (128, 256),
    upsample_strides: Sequence[float] = (1, 2),
    num_upsample_filters: Sequence[int] = (256, 256),
    shared_conv_num_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: DetectionModel

SECOND with a multi-group anchor head (nuScenes 10-class, packed point format).

Reference implementation: open-mmlab/OpenPCDet (cbgs_second_multihead). A residual sparse 3D backbone (VoxelResBackbone8x) feeds the shared 2D BEV backbone and an MultiGroupAnchorHead head. Input points carry 5 features (\(x, y, z, \text{intensity}, \Delta t\)).

Parameters:

  • in_channels (int, default: 5 ) –

    Raw point feature channels including xyz (5 for nuScenes).

  • num_classes (int, default: 10 ) –

    Number of foreground classes (10 for nuScenes).

  • voxel_size (Sequence[float], default: (0.1, 0.1, 0.2) ) –

    Voxel size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (-51.2, -51.2, -5.0, 51.2, 51.2, 3.0) ) –

    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)\), one row per class.

  • anchor_bottom_heights (Sequence[float]) –

    Per-class anchor bottom \(z\), one per class.

  • head_class_groups (Sequence[Sequence[int]]) –

    Class-index groups, one per RPN head (e.g. [[0], [1, 2], ...]).

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

  • layer_nums (Sequence[int], default: (5, 5) ) –

    2D backbone conv counts per level.

  • layer_strides (Sequence[int], default: (1, 2) ) –

    2D backbone downsample strides per level.

  • num_filters (Sequence[int], default: (128, 256) ) –

    2D backbone channel widths per level.

  • upsample_strides (Sequence[float], default: (1, 2) ) –

    2D backbone upsample factors per level.

  • num_upsample_filters (Sequence[int], default: (256, 256) ) –

    2D backbone upsample channels per level.

  • shared_conv_num_filter (int, default: 64 ) –

    Channels of the head's shared conv.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation type or callable for the 3D/2D backbones and head.

  • 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 3D/2D backbones and head.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra normalization arguments.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_backbone_3d –

    Build the residual sparse 3D voxel backbone.

  • configure_backbone –

    Build the 2D BEV backbone.

  • configure_head –

    Build the multi-group anchor head.

  • decode –

    Decode a forward output into raw per-anchor detections (see MultiGroupAnchorHead.decode).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the BEV feature map entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_backbone_3d [source] ΒΆ

configure_backbone_3d() -> VoxelResBackbone8x

Build the residual sparse 3D voxel backbone.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVBackbone

Build the 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> MultiGroupAnchorHead

Build the multi-group anchor head.

decode [source] ΒΆ

Decode a forward output into raw per-anchor detections (see MultiGroupAnchorHead.decode).

SonataSegmentation [source] ΒΆ

SonataSegmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "z",
        "z-trans",
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (3, 3, 3, 12, 3),
    encoder_channels: Sequence[int] = (
        48,
        96,
        192,
        384,
        512,
    ),
    encoder_num_heads: Sequence[int] = (3, 6, 12, 24, 32),
    encoder_patch_size: Sequence[int] = (
        1024,
        1024,
        1024,
        1024,
        1024,
    ),
    norm: Union[str, Callable] = "layer_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    dropout: float = 0.0,
    pooling: str = "grid",
    stem_type: str = "linear",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
)

Bases: SemanticSegmentationModel

Sonata linear-probing segmentation model.

This variant follows the segmentation demo from facebookresearch/sonata: the encoder features are unpooled through the saved pooling inverses, concatenated with each parent stage, then projected by a linear segmentation head.

Note

The default (and registered) configuration enables flash attention (use_flash_attn=True), which requires the flash-attn package and a CUDA device; pass use_flash_attn=False to run without it. The xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run the forward on CPU.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

forward [source] ΒΆ

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

Forward pass.

Parameters:

  • x (Tensor) –

    Per-point features of shape \((N, C)\).

  • pos_grid (Tensor) –

    Integer voxel-grid coordinates of shape \((N, 3)\) (used by the encoder for Z-order / Hilbert serialization, not float positions).

  • batch (Tensor) –

    Per-point batch index of shape \((N,)\).

SPFormerUNetSegmentation [source] ΒΆ

SPFormerUNetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    channels: Sequence[int] = (32, 64, 96, 128, 160),
    depths: Union[int, Sequence[int]] = 2,
    stem_kernel_size: int = 3,
    spatial_padding: int = 96,
    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: SemanticSegmentationModel

SpConv U-Net from SPFormer.

Reference: sunjiahao1999/SPFormer. A symmetric submanifold sparse-convolution U-Net with pre-norm residual blocks. Distinct from SparseUNetSegmentation: the residual blocks are pre-norm and the stem-level blocks run at full resolution before any downsampling.

Set num_classes=0 to drop the classifier: forward then returns the normalized per-voxel features \((N, \text{channels}[0])\).

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of output classes; 0 yields an identity head.

  • channels (Sequence[int], default: (32, 64, 96, 128, 160) ) –

    Per-level channel widths, deepest level last.

  • depths (Union[int, Sequence[int]], default: 2 ) –

    Number of residual blocks per level; an int is broadcast to every level.

  • stem_kernel_size (int, default: 3 ) –

    Kernel size of the submanifold stem convolution.

  • spatial_padding (int, default: 96 ) –

    Padding (in voxels) added to the inferred spatial shape.

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

Shape
  • Input: packed features \((N, \text{in\_channels})\), grid coordinates, batch.
  • Output: \((N, \text{num\_classes})\) logits, or \((N, \text{channels}[0])\) features when num_classes=0.

Methods:

  • configure_encoder –

    Builds the sparse encoder producing the bottleneck features and the per-stage skips.

  • configure_decoder –

    Builds the sparse decoder upsampling the bottleneck back to full resolution.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the full-resolution decoder features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the full-resolution decoder features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> SPFormerUNetEncoder

Builds the sparse encoder producing the bottleneck features and the per-stage skips.

configure_decoder [source] ΒΆ

configure_decoder() -> SPFormerUNetDecoder

Builds the sparse decoder upsampling the bottleneck back to full resolution.

SphereFormerSegmentation [source] ΒΆ

SphereFormerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: int = 32,
    channels: Sequence[int] = (32, 64, 128, 256, 256),
    depth: int = 2,
    head_dim: int = 16,
    window_size: Sequence[float] = (0.3, 0.3, 0.3),
    window_size_sphere: Sequence[float] = (2.0, 2.0, 80.0),
    quant_size: Sequence[float] = (0.0125, 0.0125, 0.0125),
    quant_size_sphere: Sequence[float] = (
        2.0 / 24,
        2.0 / 24,
        80.0 / 24,
    ),
    window_size_scale: Tuple[float, float] = (2.0, 1.5),
    sphere_layers: Sequence[int] = (1, 2, 3, 4, 5),
    radial_split_exponent: float = 0.0125,
    drop_path: float = 0.0,
    min_spatial_shape: int = 128,
    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: SemanticSegmentationModel

SphereFormer semantic-segmentation model, as described in the paper Spherical Transformer for LiDAR-based 3D Recognition.

A sparse-convolution UNet32 backbone with a cubic + radial windowed self-attention block at every stage. The windowed attention is computed by the sptr CUDA kernel (an optional dependency), as in the reference. Inputs follow the packed convention: point features x, integer voxel-grid coordinates pos_grid, real-valued coordinates pos (used by the attention), and a per-point batch index. The output is per-point class logits.

Parameters:

  • in_channels (int) –

    Input feature channels.

  • num_classes (int) –

    Number of semantic classes.

  • stem_channels (int, default: 32 ) –

    Stem / level-0 channel count \(m\).

  • channels (Sequence[int], default: (32, 64, 128, 256, 256) ) –

    Per-level channel counts (length = number of UNet levels).

  • depth (int, default: 2 ) –

    Residual blocks per level (before and after the recursive branch).

  • head_dim (int, default: 16 ) –

    Per-head dimension for the windowed attention.

  • window_size (Sequence[float], default: (0.3, 0.3, 0.3) ) –

    Base cubic window size (voxel_size * patch_size * window), of shape \((3,)\).

  • window_size_sphere (Sequence[float], default: (2.0, 2.0, 80.0) ) –

    Base spherical window size \((\theta, \phi, r)\), of shape \((3,)\).

  • quant_size (Sequence[float], default: (0.0125, 0.0125, 0.0125) ) –

    Base cubic quantization size, of shape \((3,)\).

  • quant_size_sphere (Sequence[float], default: (2.0 / 24, 2.0 / 24, 80.0 / 24) ) –

    Base spherical quantization size, of shape \((3,)\).

  • window_size_scale (Tuple[float, float], default: (2.0, 1.5) ) –

    Pair (cubic_scale, sphere_scale) applied per deeper level.

  • sphere_layers (Sequence[int], default: (1, 2, 3, 4, 5) ) –

    Levels (1-indexed) that receive a windowed-attention block.

  • radial_split_exponent (float, default: 0.0125 ) –

    Base bin width for the radial exponential split.

  • drop_path (float, default: 0.0 ) –

    Maximum stochastic-depth rate (linearly spread across levels).

  • min_spatial_shape (int, default: 128 ) –

    Per-axis lower bound on the inferred sparse spatial shape.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer name / callable.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the normalization layer.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation name / callable.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the activation.

Example
>>> import torch
>>> from torch_pointcloud.models import create_model
>>> model = create_model("sphereformer.semantickitti", task="semantic-segmentation").eval()  # doctest: +SKIP
>>> pos = torch.rand(1000, 3) * 10  # doctest: +SKIP
>>> pos_grid = (pos / 0.05).floor().long()  # doctest: +SKIP
>>> x = torch.cat([pos, torch.rand(1000, 1)], dim=1)  # doctest: +SKIP
>>> batch = torch.zeros(1000, dtype=torch.long)  # doctest: +SKIP
>>> logits = model(x, pos, pos_grid, batch)  # doctest: +SKIP
>>> logits.shape  # doctest: +SKIP
torch.Size([1000, 19])

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the full-resolution features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the full-resolution features entering the head.

configure_input_conv [source] ΒΆ

configure_input_conv() -> Module

Build the submanifold stem convolution.

configure_unet [source] ΒΆ

configure_unet() -> SphereFormerUBlock

Build the recursive sparse UNet with windowed attention at every level.

configure_output_layer [source] ΒΆ

configure_output_layer() -> Module

Build the final normalization and activation applied before the head.

SparseUNetSegmentation [source] ΒΆ

SparseUNetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: int = 32,
    encoder_channels: Sequence[int] = (32, 64, 128, 256),
    decoder_channels: Sequence[int] = (256, 128, 96, 96),
    encoder_depths: Sequence[int] = (2, 3, 4, 6),
    decoder_depths: Sequence[int] = (2, 2, 2, 2),
    stem_kernel_size: int = 5,
    kernel_size: int = 3,
    spatial_padding: int = 96,
    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: SemanticSegmentationModel

SpUNet segmentation model, a sparse residual U-Net in the spirit of 4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks by Christopher Choy, JunYoung Gwak, Silvio Savarese.

Voxelized coordinates enter a sparse convolutional stem, then a symmetric encoder-decoder of submanifold residual blocks with skip connections. The head runs on the full-resolution decoder output, whose rows stay aligned with the input points.

Parameters:

  • in_channels (int) –

    Number of input feature channels. Positions are used as features when x is None.

  • num_classes (int) –

    Number of output classes. \(0\) replaces the head with nn.Identity.

  • stem_channels (int, default: 32 ) –

    Number of channels of the stem, also the width of the shallowest skip.

  • encoder_channels (Sequence[int], default: (32, 64, 128, 256) ) –

    Feature width of every encoder stage.

  • decoder_channels (Sequence[int], default: (256, 128, 96, 96) ) –

    Feature width of every decoder stage, deepest first.

  • encoder_depths (Sequence[int], default: (2, 3, 4, 6) ) –

    Number of residual blocks per encoder stage.

  • decoder_depths (Sequence[int], default: (2, 2, 2, 2) ) –

    Number of residual blocks per decoder stage, aligned with decoder_channels.

  • stem_kernel_size (int, default: 5 ) –

    Kernel size of the stem convolution.

  • kernel_size (int, default: 3 ) –

    Kernel size of the residual convolutions.

  • spatial_padding (int, default: 96 ) –

    Padding added to the sparse spatial shape, so that voxel indices stay in range.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the activation function.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization function.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Keyword arguments for the normalization function.

Methods:

  • configure_encoder –

    Build the sparse encoder producing the bottleneck features and the per-stage skips.

  • configure_decoder –

    Build the sparse decoder upsampling the bottleneck back to full resolution.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the full-resolution decoder features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the full-resolution decoder features entering the head.

configure_encoder [source] ΒΆ

configure_encoder() -> SparseUNetEncoder

Build the sparse encoder producing the bottleneck features and the per-stage skips.

configure_decoder [source] ΒΆ

configure_decoder() -> SparseUNetDecoder

Build the sparse decoder upsampling the bottleneck back to full resolution.

SPVCNNClassification [source] ΒΆ

SPVCNNClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stem_channels: int = 32,
    encoder_channels: Sequence[int] = (32, 64, 128, 256),
    encoder_depths: Sequence[int] = (2, 2, 2, 2),
    encoder_fusion_stages: Sequence[bool] = (
        False,
        False,
        False,
        True,
    ),
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.3,
    global_pool: PoolLike = "max",
    dropout: 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: ClassificationModel

SPVCNN classification model as described in the paper Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution by Haotian Tang, Zhijian Liu, Shengyu Zhao, Yujun Lin, Ji Lin, Hanrui Wang, Song Han.

Sparse voxel convolutions carry the coarse context while a parallel point branch keeps the fine geometry, the two being merged at every fusion stage.

Parameters:

  • in_channels (int) –

    Number of input channels. Falls back to spatial_dim when None.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of the point clouds.

  • stem_channels (int, default: 32 ) –

    Number of channels of the two stem convolutions.

  • encoder_channels (Sequence[int], default: (32, 64, 128, 256) ) –

    Output channels of each encoder stage.

  • encoder_depths (Sequence[int], default: (2, 2, 2, 2) ) –

    Number of residual blocks in each encoder stage.

  • encoder_fusion_stages (Sequence[bool], default: (False, False, False, True) ) –

    Whether each encoder stage ends with a point-voxel fusion block.

  • kernel_size (int, default: 3 ) –

    Kernel size of the residual convolutions.

  • stride (int, default: 1 ) –

    Stride of the residual convolutions.

  • dilation (int, default: 1 ) –

    Dilation of the residual convolutions.

  • drop_path (float, default: 0.3 ) –

    Maximum stochastic depth rate, ramped linearly across the encoder blocks.

  • global_pool (PoolLike, default: 'max' ) –

    Global pooling used to reduce the point features to one vector per cloud.

  • dropout (float, default: 0.0 ) –

    Dropout applied to the pooled features before the head.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the activation factory.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the normalization factory.

Methods:

  • configure_stem –

    Build the two-convolution sparse stem lifting the input features to stem_channels.

  • configure_encoder –

    Build the SPVCNNEncoder backbone.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the encoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the encoder output.

configure_stem [source] ΒΆ

configure_stem() -> Sequential

Build the two-convolution sparse stem lifting the input features to stem_channels.

configure_encoder [source] ΒΆ

configure_encoder() -> SPVCNNEncoder

Build the SPVCNNEncoder backbone.

SPVCNNSegmentation [source] ΒΆ

SPVCNNSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stem_channels: int = 32,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_fusion_stages: Sequence[bool],
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_fusion_stages: Sequence[bool],
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.3,
    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: SemanticSegmentationModel

SPVCNN segmentation model as described in the paper Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution by Haotian Tang, Zhijian Liu, Shengyu Zhao, Yujun Lin, Ji Lin, Hanrui Wang, Song Han.

A sparse voxel U-Net with skip connections, run alongside a point branch that the fusion stages merge back into the voxel features. Per-point logits are read from the point branch.

Parameters:

  • in_channels (int) –

    Number of input channels. Falls back to spatial_dim when None.

  • num_classes (int) –

    Number of output classes.

  • spatial_dim (int, default: 3 ) –

    Spatial dimensionality of the point clouds.

  • stem_channels (int, default: 32 ) –

    Number of channels of the two stem convolutions.

  • encoder_channels (Sequence[int]) –

    Output channels of each encoder stage.

  • encoder_depths (Sequence[int]) –

    Number of residual blocks in each encoder stage.

  • encoder_fusion_stages (Sequence[bool]) –

    Whether each encoder stage ends with a point-voxel fusion block.

  • decoder_channels (Sequence[int]) –

    Output channels of each decoder stage.

  • decoder_depths (Sequence[int]) –

    Number of residual blocks in each decoder stage.

  • decoder_fusion_stages (Sequence[bool]) –

    Whether each decoder stage ends with a point-voxel fusion block.

  • kernel_size (int, default: 3 ) –

    Kernel size of the residual convolutions.

  • stride (int, default: 1 ) –

    Stride of the residual convolutions.

  • dilation (int, default: 1 ) –

    Dilation of the residual convolutions.

  • drop_path (float, default: 0.3 ) –

    Maximum stochastic depth rate, ramped linearly across the encoder blocks.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation function.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the activation factory.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization layer.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Optional keyword arguments for the normalization factory.

Methods:

  • configure_stem –

    Build the two-convolution sparse stem lifting the input features to stem_channels.

  • configure_encoder –

    Build the SPVCNNEncoder backbone.

  • configure_decoder –

    Build the SPVCNNDecoder upsampling the coarsest features back through the encoder skips.

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the decoder output.

num_features property ΒΆ

num_features: int

Feature dimension \(C\) of the decoder output.

configure_stem [source] ΒΆ

configure_stem() -> Sequential

Build the two-convolution sparse stem lifting the input features to stem_channels.

configure_encoder [source] ΒΆ

configure_encoder() -> SPVCNNEncoder

Build the SPVCNNEncoder backbone.

configure_decoder [source] ΒΆ

configure_decoder() -> SPVCNNDecoder

Build the SPVCNNDecoder upsampling the coarsest features back through the encoder skips.

ThreeDETRDetection [source] ΒΆ

ThreeDETRDetection(
    in_channels: int,
    num_classes: int,
    *,
    num_heading_bins: int,
    num_queries: int,
    preencoder_num_points: int = 2048,
    encoder_type: str = "vanilla",
    encoder_embed_dim: int = 256,
    encoder_num_heads: int = 4,
    encoder_feedforward_channels: int = 128,
    encoder_depth: int = 3,
    encoder_dropout: float = 0.1,
    decoder_embed_dim: int = 256,
    decoder_num_heads: int = 4,
    decoder_feedforward_channels: int = 256,
    decoder_depth: int = 8,
    decoder_dropout: float = 0.1,
    mlp_dropout: float = 0.3,
    preenc_radius: float = 0.2,
    preenc_nsample: 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: DetectionModel

3DETR end-to-end transformer 3D object detector (packed point format).

Reference: An End-to-End Transformer Model for 3D Object Detection (Misra et al., 2021). Reference implementation: facebookresearch/3detr.

A set-abstraction tokenizer downsamples the cloud to a fixed set of point tokens, a transformer encoder (vanilla, or masked with one interim downsampling for 3DETR-m) refines them, a fixed number of object queries are farthest-point-sampled from the encoder tokens, and a transformer decoder attends them against the encoder memory. Five MLP heads decode each query into a class, center, size and heading. Box centers and sizes are predicted in a per-scene min-max normalized frame and unnormalized against the point-cloud extent.

Parameters:

  • in_channels (int) –

    Input feature channels per point excluding xyz (\(0\) for xyz-only, \(3\) for RGB).

  • num_classes (int) –

    Number of semantic classes (the class head adds one background slot).

  • num_heading_bins (int) –

    Heading-angle bins (\(1\) for axis-aligned ScanNet, \(12\) for oriented SUN RGB-D).

  • num_queries (int) –

    Number of object queries (decoded boxes) per scene.

  • preencoder_num_points (int, default: 2048 ) –

    Token count after the set-abstraction tokenizer.

  • encoder_type (str, default: 'vanilla' ) –

    "vanilla" (encoder keeps preencoder_num_points tokens) or "masked" (3DETR-m: radius attention masks plus one interim downsampling to \(\text{preenc\_npoints} // 2\)).

  • encoder_embed_dim (int, default: 256 ) –

    Encoder token dimension.

  • encoder_num_heads (int, default: 4 ) –

    Encoder attention heads.

  • encoder_feedforward_channels (int, default: 128 ) –

    Encoder feed-forward width.

  • encoder_depth (int, default: 3 ) –

    Encoder layers.

  • encoder_dropout (float, default: 0.1 ) –

    Encoder dropout.

  • decoder_embed_dim (int, default: 256 ) –

    Decoder token dimension.

  • decoder_num_heads (int, default: 4 ) –

    Decoder attention heads.

  • decoder_feedforward_channels (int, default: 256 ) –

    Decoder feed-forward width.

  • decoder_depth (int, default: 8 ) –

    Decoder layers.

  • decoder_dropout (float, default: 0.1 ) –

    Decoder dropout.

  • mlp_dropout (float, default: 0.3 ) –

    Dropout inside each box head.

  • preenc_radius (float, default: 0.2 ) –

    Tokenizer ball-query radius.

  • preenc_nsample (int, default: 64 ) –

    Tokenizer ball-query neighbor cap.

  • 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 for the convolutional blocks (tokenizer, projection, heads).

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra normalization arguments.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the encoder tokens entering the decoder.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the encoder tokens entering the decoder.

configure_pre_encoder [source] ΒΆ

configure_pre_encoder() -> PointnetSAModuleVotes

Build the set-abstraction tokenizer.

configure_encoder [source] ΒΆ

configure_encoder() -> Module

Build the transformer encoder (masked with one interim downsampling for 3DETR-m).

configure_encoder_to_decoder_projection [source] ΒΆ

configure_encoder_to_decoder_projection() -> GenericConvMLP

Build the projection from encoder tokens to the decoder dimension.

configure_pos_embedding [source] ΒΆ

configure_pos_embedding() -> PositionEmbeddingFourier

Build the Fourier position embedding.

configure_query_projection [source] ΒΆ

configure_query_projection() -> GenericConvMLP

Build the query projection.

configure_decoder [source] ΒΆ

configure_decoder() -> TransformerDecoder

Build the transformer decoder.

configure_mlp_heads [source] ΒΆ

configure_mlp_heads() -> ModuleDict

Build the per-query class, center, size and heading heads.

run_encoder [source] ΒΆ

run_encoder(
    x: OptTensor,
    pos: Tensor,
    batch: Tensor,
    idx: OptTensor = None,
) -> Tuple[Tensor, Tensor]

Tokenizes the point cloud with the pre-encoder and runs the masked transformer encoder.

Parameters:

  • x (OptTensor) –

    Packed point features, shape \((N, C)\), or None to use the positions.

  • pos (Tensor) –

    Packed point positions, shape \((N, 3)\).

  • batch (Tensor) –

    Per-point scene index, shape \((N,)\).

  • idx (OptTensor, default: None ) –

    Optional pre-computed sampling indices for the pre-encoder.

Returns:

  • Tuple[Tensor, Tensor] –

    The token positions of shape \((B, P, 3)\) and the token features of shape \((P, B, C)\).

get_query_embeddings [source] ΒΆ

get_query_embeddings(
    enc_xyz: Tensor,
    point_cloud_dims: Tuple[Tensor, Tensor],
    query_idx: OptTensor = None,
) -> Tuple[Tensor, Tensor]

Samples the query positions among the encoder tokens and embeds them into decoder queries.

Parameters:

  • enc_xyz (Tensor) –

    Encoder token positions, shape \((B, P, 3)\).

  • point_cloud_dims (Tuple[Tensor, Tensor]) –

    Per-scene minimum and maximum corners, used to normalize the positional embedding.

  • query_idx (OptTensor, default: None ) –

    Optional pre-computed query indices. Farthest point sampling is used when None.

Returns:

  • Tuple[Tensor, Tensor] –

    The query positions of shape \((B, Q, 3)\) and the query embeddings of shape \((B, C, Q)\).

decode [source] ΒΆ

decode(out: ThreeDETROutput) -> Detection3D

Decode a forward output into raw per-query detections (no NMS, threshold, or filtering).

Builds one oriented box per query, scores it by objectness, and labels it by the argmax semantic class. The angle head predicts negated angles, so the decoded heading is the negated angle_continuous, matching the dataset / metric \((c_x, c_y, c_z, d_x, d_y, d_z, \theta)\) counter-clockwise convention. The result is the full unfiltered query set; the evaluation pipeline applies point-count filtering, NMS, score thresholding, and the indoor per-class expansion (driven by the returned class_probs) via the torch_pointcloud.ops.box3d utilities, reproducing 3DETR's APCalculator test protocol (exact_eval=True).

Parameters:

Returns:

  • Detection3D –

    Packed queries {"boxes", "scores", "labels", "batch", "class_probs"} (PyG layout), where the

  • Detection3D –

    per-query score is objectness, the label is the argmax semantic class, and class_probs holds

  • Detection3D –

    the semantic-class probabilities.

Shape
  • boxes: \((B \cdot Q, 7)\)
  • scores / labels / batch: \((B \cdot Q,)\)
  • class_probs: \((B \cdot Q, C)\)

UtoniaSegmentation [source] ΒΆ

UtoniaSegmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "z",
        "z-trans",
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (3, 3, 3, 12, 3),
    encoder_channels: Sequence[int] = (
        54,
        108,
        216,
        432,
        576,
    ),
    encoder_num_heads: Sequence[int] = (3, 6, 12, 24, 32),
    encoder_patch_size: Sequence[int] = (
        1024,
        1024,
        1024,
        1024,
        1024,
    ),
    norm: Union[str, Callable] = "layer_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    rope_base: float = 10.0,
    dropout: float = 0.0,
    pooling: str = "grid",
    stem_type: str = "linear",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
)

Bases: SemanticSegmentationModel

Utonia linear-probing segmentation model.

Linear-probe variant from Utonia: Toward One Encoder for All Point Clouds (ICML 2026). Architecturally similar to Sonata / Concerto's linear-probe head, with one key change: every attention layer adds a 3D rotary position embedding (Point3DRoPE) on top of (q, k), indexed by the real-valued metric position rather than the integer voxel grid. The position is mean-pooled at every encoder stage so each level operates at its natural scale.

Note

The default (and registered) configuration enables flash attention (use_flash_attn=True), which requires the flash-attn package and a CUDA device; pass use_flash_attn=False to run without it. The xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run the forward on CPU.

Methods:

  • configure_encoder –

    Build the PointTransformerV3Encoder backbone with rotary position embeddings.

  • forward –

    Forward pass.

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

num_features property ΒΆ

num_features: int

Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

configure_encoder [source] ΒΆ

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone with rotary position embeddings.

forward [source] ΒΆ

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

Forward pass.

Parameters:

  • x (Tensor) –

    Per-point features of shape \((N, C)\).

  • pos (Tensor) –

    Real-valued metric positions of shape \((N, 3)\) used by 3D RoPE.

  • pos_grid (Tensor) –

    Integer voxel-grid coordinates of shape \((N, 3)\) used by the encoder for serialization and sparse convolutions.

  • batch (Tensor) –

    Per-point batch index of shape \((N,)\).

VoteNetDetection [source] ΒΆ

VoteNetDetection(
    in_channels: int,
    num_classes: int,
    *,
    num_heading_bins: int,
    num_size_clusters: int,
    mean_sizes: Union[Tensor, List[List[float]]],
    num_proposals: int,
    vote_factor: int,
    sampling: str,
    sa_channels: Sequence[Sequence[int]],
    sa_num_points: Sequence[int],
    sa_radii: Sequence[float],
    sa_num_neighbors: Sequence[int],
    fp_channels: Sequence[Sequence[int]],
    vote_aggr_channels: Sequence[int],
    vote_aggr_radius: float,
    vote_aggr_num_neighbors: 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: DetectionModel

VoteNet 3D object detector (packed point format).

Reference: Deep Hough Voting for 3D Object Detection in Point Clouds (Qi et al., 2019). Reference implementation: facebookresearch/votenet.

A PointNet++ backbone extracts seed points; a voting module shifts each seed toward its object center; a proposal module clusters the votes and decodes a fixed number of oriented (or axis-aligned) box proposals per scene.

Parameters:

  • in_channels (int) –

    Input feature channels per point excluding xyz (e.g. \(1\) for a floor-relative height feature, \(4\) for height + RGB).

  • num_classes (int) –

    Number of semantic classes.

  • num_heading_bins (int) –

    Number of heading-angle bins (\(1\) for axis-aligned ScanNet boxes, \(12\) for oriented SUN RGB-D boxes).

  • num_size_clusters (int) –

    Number of size templates (one per class here).

  • mean_sizes (Union[Tensor, List[List[float]]]) –

    Per-template mean box size, shape \((\text{num\_size\_cluster}, 3)\).

  • num_proposals (int) –

    Number of box proposals per scene.

  • vote_factor (int) –

    Votes generated per seed.

  • sampling (str) –

    Aggregation-center sampling, "vote_fps" or "seed_fps".

  • sa_channels (Sequence[Sequence[int]]) –

    Per-SA-block MLP channel lists for the backbone.

  • sa_num_points (Sequence[int]) –

    Per-SA-block farthest-point-sample counts.

  • sa_radii (Sequence[float]) –

    Per-SA-block ball-query radii.

  • sa_num_neighbors (Sequence[int]) –

    Per-SA-block neighbor caps.

  • fp_channels (Sequence[Sequence[int]]) –

    Per-FP-block MLP channel lists for the backbone.

  • vote_aggr_channels (Sequence[int]) –

    MLP channels of the proposal vote-aggregation layer.

  • vote_aggr_radius (float) –

    Ball-query radius of the vote-aggregation layer.

  • vote_aggr_num_neighbors (int) –

    Neighbor cap of the vote-aggregation 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.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) of the per-seed backbone features entering the voting head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the per-seed backbone features entering the voting head.

configure_backbone [source] ΒΆ

configure_backbone() -> VoteNetBackbone

Build the PointNet++ seed backbone.

configure_vgen [source] ΒΆ

configure_vgen() -> VotingModule

Build the voting module.

configure_proposal [source] ΒΆ

configure_proposal() -> VoteNetProposalModule

Build the proposal module.

decode [source] ΒΆ

decode(out: VoteNetOutput) -> Detection3D

Decode a forward output into raw per-proposal detections (no NMS, threshold, or filtering).

Builds one oriented box per proposal from the predicted heading/size bins, scores it by objectness, and labels it by the argmax semantic class. The heading head predicts negated angles, so the decoded heading is negated to return counter-clockwise headings (the library box convention). The result is the full unfiltered proposal set; the evaluation pipeline applies point-count filtering, NMS, score thresholding, and the indoor per-class expansion (driven by the returned class_probs) via the torch_pointcloud.ops.box3d utilities.

Parameters:

Returns:

  • Detection3D –

    Packed proposals {"boxes", "scores", "labels", "batch", "class_probs"} (PyG layout), where the

  • Detection3D –

    per-proposal score is objectness, the label is the argmax semantic class, and class_probs holds

  • Detection3D –

    the softmaxed semantic-class probabilities.

Shape
  • boxes: \((B \cdot P, 7)\)
  • scores / labels / batch: \((B \cdot P,)\)
  • class_probs: \((B \cdot P, C)\)

VoxelMambaDetection [source] ΒΆ

VoxelMambaDetection(
    in_channels: int = 5,
    num_classes: int = 3,
    *,
    voxel_size: Sequence[float] = (0.32, 0.32, 0.1875),
    point_cloud_range: Sequence[float] = (
        -74.88,
        -74.88,
        -2.0,
        74.88,
        74.88,
        4.0,
    ),
    embed_dim: int = 128,
    vfe_num_filters: Sequence[int] = (128, 128),
    layer_nums: Sequence[int] = (1, 2, 2),
    layer_strides: Sequence[int] = (1, 2, 2),
    num_filters: Sequence[int] = (128, 128, 256),
    upsample_strides: Sequence[float] = (1, 2, 4),
    num_upsample_filters: Sequence[int] = (128, 128, 128),
    shared_conv_channels: int = 64,
    rms_norm: bool = True,
    fused_add_norm: bool = True,
    norm_epsilon: float = 1e-05,
)

Bases: DetectionModel

Voxel Mamba: group-free state-space 3D object detector (packed point format).

Reference: Zhang et al., 2024. Reference implementation: gwenzhang/Voxel-Mamba (built on DSVT).

Voxels are serialized into a single Hilbert-curve sequence and processed by bidirectional Mamba (state-space) blocks (no windowing / grouping), then scattered to a BEV map, refined by a 2D residual backbone, and decoded by a center-based head.

Parameters:

  • in_channels (int, default: 5 ) –

    Raw point feature channels including xyz (e.g. \(5\) for Waymo).

  • num_classes (int, default: 3 ) –

    Number of foreground classes.

  • voxel_size (Sequence[float], default: (0.32, 0.32, 0.1875) ) –

    Voxel size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (-74.88, -74.88, -2.0, 74.88, 74.88, 4.0) ) –

    Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).

  • embed_dim (int, default: 128 ) –

    Voxel feature channels of the Mamba backbone.

  • vfe_num_filters (Sequence[int], default: (128, 128) ) –

    PFN widths of the dynamic mean VFE.

  • layer_nums (Sequence[int], default: (1, 2, 2) ) –

    2D backbone residual-block counts per level.

  • layer_strides (Sequence[int], default: (1, 2, 2) ) –

    2D backbone downsample strides per level.

  • num_filters (Sequence[int], default: (128, 128, 256) ) –

    2D backbone channel widths per level.

  • upsample_strides (Sequence[float], default: (1, 2, 4) ) –

    2D backbone upsample factors per level.

  • num_upsample_filters (Sequence[int], default: (128, 128, 128) ) –

    2D backbone upsample channels per level.

  • shared_conv_channels (int, default: 64 ) –

    Channels of the head's shared conv.

  • rms_norm (bool, default: True ) –

    Use RMSNorm inside the Mamba blocks.

  • fused_add_norm (bool, default: True ) –

    Use the fused add+norm kernel inside the Mamba blocks.

  • norm_epsilon (float, default: 1e-05 ) –

    LayerNorm epsilon for the Mamba output norms.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_vfe –

    Build the dynamic mean voxel feature encoder.

  • configure_backbone_3d –

    Build the Hilbert-serialized Mamba voxel backbone.

  • configure_backbone –

    Build the residual 2D BEV backbone.

  • configure_head –

    Build the center-based detection head.

  • decode –

    Decode center-head predictions into raw candidate detections (no NMS).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the BEV feature map entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_vfe [source] ΒΆ

configure_vfe() -> DynamicMeanVFE

Build the dynamic mean voxel feature encoder.

configure_backbone_3d [source] ΒΆ

configure_backbone_3d() -> VoxelMambaBackbone

Build the Hilbert-serialized Mamba voxel backbone.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVResidualBackbone

Build the residual 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> CenterHead

Build the center-based detection head.

decode [source] ΒΆ

decode(
    out: CenterHeadOutput,
    *,
    score_threshold: float = 0.0,
    top_k: int = 500,
    iou_rectifier: Sequence[float] = (0.68, 0.71, 0.65),
) -> Detection3D

Decode center-head predictions into raw candidate detections (no NMS).

Peaks of the (sigmoid) heatmap give candidate centers; box attributes are gathered at those peaks, mapped to world coordinates, and rescored by the predicted IoU (\(s^{1 - r_c} \cdot \text{iou}^{r_c}\) with a per-class rectifier \(r_c\), as in the reference). The full candidate set is returned; the evaluation pipeline applies score thresholding and per-class 3D NMS via the torch_pointcloud.ops.box3d utilities.

Parameters:

  • out (CenterHeadOutput) –

    A CenterHeadOutput from forward.

  • score_threshold (float, default: 0.0 ) –

    Minimum (pre-rectification) heatmap score to keep a peak; the non-filtering \(0\) default returns every peak (the reference protocol filters at \(0.1\)).

  • top_k (int, default: 500 ) –

    Number of heatmap peaks gathered per scene.

  • iou_rectifier (Sequence[float], default: (0.68, 0.71, 0.65) ) –

    Per-class IoU-rectification exponent, one entry per class (the default holds the reference Waymo 3-class values).

Returns:

  • Detection3D –

    Packed candidate detections {"boxes": (K, 7), "scores": (K,), "labels": (K,), "batch": (K,)}

  • Detection3D –

    (PyG layout).

VoxelNeXtDetection [source] ΒΆ

VoxelNeXtDetection(
    in_channels: int = 5,
    num_classes: int = 10,
    *,
    voxel_size: Sequence[float] = (0.075, 0.075, 0.2),
    point_cloud_range: Sequence[float] = (
        -54.0,
        -54.0,
        -5.0,
        54.0,
        54.0,
        3.0,
    ),
    head_class_groups: Sequence[Sequence[int]],
    feature_map_stride: int,
    channels: Sequence[int] = (16, 32, 64, 128, 128),
    shared_conv_channels: int = 128,
    head_kernel_size: int = 1,
    num_hm_conv: int = 2,
    use_bias_before_norm: 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: DetectionModel

VoxelNeXt fully sparse 3D object detector (packed point format).

Reference: VoxelNeXt: Fully Sparse VoxelNet for 3D Object Detection and Tracking (Chen et al., 2023). Reference implementation: dvlab-research/VoxelNeXt (ported via open-mmlab/OpenPCDet, cbgs_voxel0075_voxelnext). A residual sparse 3D backbone (VoxelResBackbone8xVoxelNeXt) collapses to a 2D BEV sparse tensor that a fully sparse multi-group head (VoxelNeXtHead) predicts boxes on directly, with no dense bird's-eye-view map. Input points carry 5 features (\(x, y, z, \text{intensity}, \Delta t\)).

Parameters:

  • in_channels (int, default: 5 ) –

    Raw point feature channels including xyz (5 for nuScenes).

  • num_classes (int, default: 10 ) –

    Number of foreground classes (10 for nuScenes).

  • voxel_size (Sequence[float], default: (0.075, 0.075, 0.2) ) –

    Voxel size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (-54.0, -54.0, -5.0, 54.0, 54.0, 3.0) ) –

    Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).

  • head_class_groups (Sequence[Sequence[int]]) –

    Class-index groups, one per separate head (e.g. [[0], [1, 2], ...]).

  • feature_map_stride (int) –

    BEV feature-map stride of the head.

  • channels (Sequence[int], default: (16, 32, 64, 128, 128) ) –

    Per-stage channel widths of the 3D backbone.

  • shared_conv_channels (int, default: 128 ) –

    Output channels of the backbone's 2D shared conv (head input).

  • head_kernel_size (int, default: 1 ) –

    Kernel size of the hidden head convs.

  • num_hm_conv (int, default: 2 ) –

    Number of convs in the classification (hm) stack.

  • use_bias_before_norm (bool, default: True ) –

    Whether hidden head convs carry a bias before their norm.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation type or callable for the backbone and head.

  • 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 backbone and head.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra normalization arguments.

Methods:

  • reset_classifier –

    Replace the classification branch of the detection head for num_classes outputs.

  • configure_backbone_3d –

    Build the fully sparse residual voxel backbone.

  • configure_head –

    Build the fully sparse multi-group detection head.

  • decode –

    Decode a forward output into raw candidate detections (see VoxelNeXtHead.decode).

Attributes:

  • num_features (int) –

    Channel count \(C\) of the sparse voxel features entering the head.

num_features property ΒΆ

num_features: int

Channel count \(C\) of the sparse voxel features entering the head.

reset_classifier [source] ΒΆ

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

configure_backbone_3d [source] ΒΆ

configure_backbone_3d() -> VoxelResBackbone8xVoxelNeXt

Build the fully sparse residual voxel backbone.

configure_head [source] ΒΆ

configure_head() -> VoxelNeXtHead

Build the fully sparse multi-group detection head.

decode [source] ΒΆ

decode(
    out: VoxelNeXtHeadOutput, *, top_k: int = 500
) -> Detection3D

Decode a forward output into raw candidate detections (see VoxelNeXtHead.decode).

create_model [source] ΒΆ

create_model(
    name: str,
    task: Literal["pretraining"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[True],
    **kwargs: Any,
) -> Tuple[Module, ModelInfoDict]
create_model(
    name: str,
    task: Literal["pretraining"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[False] = False,
    **kwargs: Any,
) -> Module
create_model(
    name: str,
    task: Literal["classification"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[True],
    **kwargs: Any,
) -> Tuple[ClassificationModel, ModelInfoDict]
create_model(
    name: str,
    task: Literal["classification"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[False] = False,
    **kwargs: Any,
) -> ClassificationModel
create_model(
    name: str,
    task: Literal["semantic-segmentation"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[True],
    **kwargs: Any,
) -> Tuple[SemanticSegmentationModel, ModelInfoDict]
create_model(
    name: str,
    task: Literal["semantic-segmentation"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[False] = False,
    **kwargs: Any,
) -> SemanticSegmentationModel
create_model(
    name: str,
    task: Literal["part-segmentation"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[True],
    **kwargs: Any,
) -> Tuple[PartSegmentationModel, ModelInfoDict]
create_model(
    name: str,
    task: Literal["part-segmentation"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[False] = False,
    **kwargs: Any,
) -> PartSegmentationModel
create_model(
    name: str,
    task: Literal["detection"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[True],
    **kwargs: Any,
) -> Tuple[DetectionModel, ModelInfoDict]
create_model(
    name: str,
    task: Literal["detection"],
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[False] = False,
    **kwargs: Any,
) -> DetectionModel
create_model(
    name: str,
    task: None = None,
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[True],
    **kwargs: Any,
) -> Tuple[Module, ModelInfoDict]
create_model(
    name: str,
    task: None = None,
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: Literal[False] = False,
    **kwargs: Any,
) -> Module
create_model(
    name: str,
    task: Optional[Task] = None,
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: bool = False,
    **kwargs: Any,
) -> Any
create_model(
    name: str,
    task: Optional[Task] = None,
    *,
    pretrained: bool = False,
    checkpoint_path: Optional[PathLike] = None,
    return_info: bool = False,
    **kwargs: Any,
) -> Any

Build a registered model, optionally loading pretrained or local checkpoint weights.

Registered weights load into the registered configuration, then an overridden num_classes rebuilds the head at its fresh initialization (0 drops it). Other overrides (e.g. in_channels) keep the matching weights, and the keys whose shape changed keep their initialization. Model keys missing from the checkpoint raise.

Parameters:

  • name (str) –

    Registered model name (see list_models).

  • task (Optional[Task], default: None ) –

    Registry the model belongs to (pretraining, classification, semantic-segmentation, part-segmentation or detection). Optional when name is registered under a single task.

  • pretrained (bool, default: False ) –

    Load the registered pretrained weights, warning when their license is non-commercial. Mutually exclusive with checkpoint_path.

  • checkpoint_path (Optional[PathLike], default: None ) –

    Local checkpoint to load instead of the registered weights. Supports torch.save files, .safetensors, and Lightning checkpoints (the wrapped network is extracted).

  • return_info (bool, default: False ) –

    Also return the registry entry (ModelInfoDict), with the resolved task and the effective hparams.

  • **kwargs (Any, default: {} ) –

    Overrides merged into the registered hparams and passed to the model constructor.

Returns:

  • Any –

    The model, or a (model, info) tuple when return_info is true.

Raises:

  • ValueError –

    If task or name is unknown, if task is omitted for a name registered under several tasks, if both pretrained and checkpoint_path are passed, or if pretrained is requested for an entry that registers no weights.

  • TypeError –

    If the entry registers architecture hparams only and the data-dependent arguments (typically in_channels and num_classes) are not passed.

  • ImportError –

    If the model needs an optional dependency that is not installed or does not meet its version requirement; the original error is chained.

  • FileNotFoundError –

    If the checkpoint file does not exist, or if the registered weights are absent from the models cache and cannot be downloaded.

Building a registered model, overriding its registered hparams, and inspecting its registry entry:

import torch_pointcloud as tp

model = tp.create_model("pointnet.modelnet40")
model = tp.create_model("pointnet.modelnet40", num_classes=10)
model, info = tp.create_model("pointnet.modelnet40", return_info=True)

deprecated_model [source] ΒΆ

deprecated_model(
    name: str,
    *,
    deprecated_in: str,
    removed_in: Optional[str] = None,
    message: Optional[str] = None,
) -> Callable[[F], F]

Keep a renamed model reachable under its former name, placed above its register_model decorator.

create_model and model_info resolve name to the registered name under its task with a FutureWarning that names both releases. list_models never lists former names, and a model of the same name registered under another task is not affected. Stack one decorator per former name, and delete it in the removed_in release.

Parameters:

  • name (str) –

    Former registry name of the entry, under the task it is registered for.

  • deprecated_in (str) –

    Release that renamed the model (e.g. 0.0.17).

  • removed_in (Optional[str], default: None ) –

    Release that drops the former name, named in the warning; None leaves it unscheduled.

  • message (Optional[str], default: None ) –

    Sentence appended to the warning, e.g. what else changed with the rename.

Returns:

  • Callable[[F], F] –

    The decorator adding the former name to the registry entry of its target callable.

Raises:

  • ValueError –

    If deprecated_in or removed_in is not a version string, if removed_in is not a later release than deprecated_in, if the decorated callable is not registered under exactly one name (the decorator belongs above register_model), or if name is the registered name itself, a registered model of its task or already a former name under its task.

Renaming a registered model while keeping its former name loadable for one release:

from torch_pointcloud.models import PointNetClassification, deprecated_model, register_model


@deprecated_model("pointnet.scanobjectnn", deprecated_in="0.0.17", removed_in="0.0.18")
@register_model("pointnet-base.scanobjectnn", task="classification", hparams=dict(in_channels=0, num_classes=15))
def pointnet_base(**hparams):
    return PointNetClassification(**hparams)

list_models [source] ΒΆ

list_models(
    name: str = "*",
    *,
    task: Optional[Task] = None,
    pretrained: bool = False,
) -> List[str]

List registered model names, sorted alphabetically.

Parameters:

  • name (str, default: '*' ) –

    Wildcard filter on the registered name (fnmatch syntax, e.g. "pointnext*").

  • task (Optional[Task], default: None ) –

    Restrict the listing to one registry; None lists across all tasks (duplicates removed).

  • pretrained (bool, default: False ) –

    Only list models that ship pretrained weights.

Returns:

  • List[str] –

    The matching model names.

Filtering by name pattern, task, and weight availability:

import torch_pointcloud as tp

classifiers = tp.list_models("pointnext*", task="classification")
with_weights = tp.list_models(task="semantic-segmentation", pretrained=True)
everything = tp.list_models()

model_info [source] ΒΆ

model_info(
    name: str, task: Optional[Task] = None
) -> ModelInfoDict

The registry entry of a model without building it: its task, weights, transform, input keys, hparams and former names.

Parameters:

  • name (str) –

    Registered model name (see list_models).

  • task (Optional[Task], default: None ) –

    Registry the model belongs to; optional when name is registered under a single task.

Returns:

  • ModelInfoDict –

    The entry as a ModelInfoDict, with the registered hparams.

Raises:

  • ValueError –

    If task or name is unknown, or if task is omitted for a name registered under several tasks.

Example
>>> info = model_info("pointnet.modelnet40")
>>> info["task"], info["input_keys"]
('classification', ('x', 'pos', 'batch'))

register_model [source] ΒΆ

register_model(
    name: str,
    *,
    hparams: Optional[Dict[str, Any]] = None,
    transform: Optional[Transform] = None,
    input_keys: Sequence[str] = ("x", "pos", "batch"),
    weights: Union[str, WeightsDict, None] = None,
    task: Literal["pretraining"],
) -> Callable[
    [Callable[..., Module]], Callable[..., Module]
]
register_model(
    name: str,
    *,
    hparams: Optional[Dict[str, Any]] = None,
    transform: Optional[Transform] = None,
    input_keys: Sequence[str] = ("x", "pos", "batch"),
    weights: Union[str, WeightsDict, None] = None,
    task: Literal["classification"],
) -> Callable[
    [Callable[..., ClassificationModel]],
    Callable[..., ClassificationModel],
]
register_model(
    name: str,
    *,
    hparams: Optional[Dict[str, Any]] = None,
    transform: Optional[Transform] = None,
    input_keys: Sequence[str] = ("x", "pos", "batch"),
    weights: Union[str, WeightsDict, None] = None,
    task: Literal["semantic-segmentation"],
) -> Callable[
    [Callable[..., SemanticSegmentationModel]],
    Callable[..., SemanticSegmentationModel],
]
register_model(
    name: str,
    *,
    hparams: Optional[Dict[str, Any]] = None,
    transform: Optional[Transform] = None,
    input_keys: Sequence[str] = ("x", "pos", "batch"),
    weights: Union[str, WeightsDict, None] = None,
    task: Literal["part-segmentation"],
) -> Callable[
    [Callable[..., PartSegmentationModel]],
    Callable[..., PartSegmentationModel],
]
register_model(
    name: str,
    *,
    hparams: Optional[Dict[str, Any]] = None,
    transform: Optional[Transform] = None,
    input_keys: Sequence[str] = ("x", "pos", "batch"),
    weights: Union[str, WeightsDict, None] = None,
    task: Literal["detection"],
) -> Callable[
    [Callable[..., DetectionModel]],
    Callable[..., DetectionModel],
]
register_model(
    name: str,
    *,
    task: Task,
    hparams: Optional[Dict[str, Any]] = None,
    transform: Optional[Transform] = None,
    input_keys: Sequence[str] = ("x", "pos", "batch"),
    weights: Union[str, WeightsDict, None] = None,
) -> Callable

Register a model entry point under name for task.

The decorated callable becomes reachable through create_model. A bare weights URL string is normalized to a WeightsDict, so registry consumers always see the structured form. A deprecated_model decorator placed above adds the former names of the entry.

Parameters:

  • name (str) –

    Registry name, <architecture>[.<dataset tag>] (e.g. pointnext-sm.scanobjectnn-hardest.openpoints).

  • task (Task) –

    Registry the model belongs to (pretraining, classification, semantic-segmentation, part-segmentation or detection).

  • hparams (Optional[Dict[str, Any]], default: None ) –

    Default keyword arguments the entry point is called with; create_model kwargs override them.

  • transform (Optional[Transform], default: None ) –

    Evaluation transform reproducing the preprocessing the weights were trained with.

  • input_keys (Sequence[str], default: ('x', 'pos', 'batch') ) –

    Batch keys forward takes positionally, in order; a dotted key reads an attribute (e.g. octree.depth). utils.data.select_inputs gathers them from a collated batch.

  • weights (Union[str, WeightsDict, None], default: None ) –

    Pretrained checkpoint, either a URL string or a WeightsDict with metadata.

Returns:

  • Callable –

    The decorator registering its target callable.

Raises:

  • ValueError –

    If task is unknown, or if name is the former name of another model of task.

Registering a classification entry point with weight metadata:

from torch_pointcloud.models import PointNetClassification, WeightsDict, register_model


@register_model(
    "pointnet-demo.scanobjectnn",
    task="classification",
    hparams=dict(in_channels=0, num_classes=15),
    weights=WeightsDict(
        url="hf://my-org/pointnet-demo.scanobjectnn/resolve/main/model.safetensors",
        dataset="scanobjectnn",
        metrics={"OA": 88.20},
    ),
)
def pointnet_demo(**hparams):
    return PointNetClassification(**hparams)