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:
-
ClassificationModelβBase class for point cloud classification models.
-
DetectionModelβBase class for 3D object detection models.
-
PartSegmentationModelβBase class for part segmentation models, registered under the
"part-segmentation"task. -
PretrainingModelβBase class for self-supervised pretraining models, registered under the
"pretraining"task. -
SemanticSegmentationModelβBase class for point cloud semantic segmentation models, registered under the
"semantic-segmentation"task. -
DeprecatedModelDictβFormer name of a registered model, declared with
deprecated_model. -
ModelDictβRegistry entry of a registered model: its entry point and the metadata
create_modelreturns. -
ModelInfoDictβRegistry entry of a built model, returned by
create_model(..., return_info=True). -
WeightsDictβStructured description of a pretrained checkpoint.
-
ConcertoSegmentationβConcerto linear-probing segmentation model.
-
DGCNNClassificationβClassification model as described in the paper
-
DGCNNPartSegmentationβPart segmentation model as described in the paper
-
DGCNNSegmentationβSemantic segmentation model as described in the paper
-
KPFCNNClassificationβKPConv Network for classification tasks as described in the paper
-
KPFCNNSegmentationβKPConv Network for segmentation tasks as described in the paper
-
LIONDetectionβLION: linear group RNN (Mamba) 3D object detector (packed point format).
-
OctFormerClassificationβOctFormer classification model from
-
OctFormerSegmentationβOctFormer segmentation model from
-
PointBERTClassificationβPoint-BERT classification model.
-
PointBERTDiscreteVAEβPoint-BERT discrete VAE point tokenizer.
-
PointBERTPretrainingβPoint-BERT masked point modeling backbone (pretrain).
-
PointM2AEClassificationβImplementation of the Point-M2AE classification model.
-
PointM2AEPartSegmentationβImplementation of the Point-M2AE part-segmentation model.
-
PointM2AEPretrainingβImplementation of the Point-M2AE pre-training model.
-
PointMAEClassificationβPoint-MAE classification model, as in [Masked Autoencoders for Point Cloud Self-supervised
-
PointMAEPartSegmentationβPoint-MAE part-segmentation model, as in [Masked Autoencoders for Point Cloud Self-supervised
-
PointMAEPretrainingβPoint-MAE masked-autoencoder pretraining model, as in [Masked Autoencoders for Point Cloud
-
PointMambaClassificationβ -
PointMambaPretrainingβPointMamba masked-autoencoder pre-training model, as described in the paper PointMamba: A Simple State Space Model for Point Cloud Analysis.
-
PointTransformerClassificationβPoint Transformer classification model from the paper
-
PointTransformerSegmentationβPoint Transformer segmentation model from the paper
-
PointTransformerV2ClassificationβImplementation of the Point Transformer V2 model for classification as described in the paper
-
PointTransformerV2SegmentationβImplementation of the Point Transformer V2 model for semantic segmentation as described in the paper
-
PointTransformerV3ClassificationβPyTorch implementation of the Point Transformer V3 model, as described in the paper
-
PointTransformerV3SegmentationβPyTorch implementation of the Point Transformer V3 model for segmentation tasks.
-
PointCNNClassificationβClassification model as described in the paper
-
PointCNNSegmentationβSegmentation model as described in the paper
-
PointConvDensityClassificationβPointConv classification model with density re-weighting from
-
PointGPTClassificationβPointGPT classification model.
-
PointGPTPretrainingβPointGPT auto-regressive generative pretraining model.
-
PointMLPClassificationβPointMLP classification model from
-
PointMLPSegmentationβPointMLP segmentation model from
-
PointNetClassificationβPointNet architecture for 3D point cloud classification tasks as described in the original PointNet paper
-
PointNetSegmentationβPointNet architecture for point cloud segmentation tasks as described in the original PointNet paper
-
PointNet2ClassificationβPointNet++ classification model from the paper
-
PointNet2SegmentationβPointNet++ segmentation model from the paper
-
PointNeXtClassificationβPointNeXt classification model as described in the paper
-
PointNeXtPartSegmentationβPointNeXt part segmentation model for ShapeNetPart.
-
PointNeXtSegmentationβPointNeXt segmentation model as described in the paper
-
PointPillarsDetectionβPointPillars 3D object detector (packed point format).
-
PointPillarsMultiHeadDetectionβPointPillars with a multi-group anchor head (nuScenes 10-class, packed point format).
-
PointRCNNDetectionβPointRCNN two-stage point-based 3D object detector (packed point format).
-
PVCNNClassificationβPVCNN classification model from
-
PVCNNSegmentationβPVCNN segmentation model from
-
PVCNN2ClassificationβPVCNN++ classification model from
-
PVCNN2SegmentationβPVCNN++ segmentation model from
-
RandLANetClassificationβRandLA-Net classification model from
-
RandLANetSegmentationβRandLA-Net segmentation model from
-
SECONDDetectionβSECOND 3D object detector (packed point format).
-
SECONDMultiHeadDetectionβSECOND with a multi-group anchor head (nuScenes 10-class, packed point format).
-
SonataSegmentationβSonata linear-probing segmentation model.
-
SPFormerUNetSegmentationβSpConv U-Net from SPFormer.
-
SphereFormerSegmentationβSphereFormer semantic-segmentation model, as described in the paper
-
SparseUNetSegmentationβSpUNet segmentation model, a sparse residual U-Net in the spirit of
-
SPVCNNClassificationβSPVCNN classification model as described in the paper
-
SPVCNNSegmentationβSPVCNN segmentation model as described in the paper
-
ThreeDETRDetectionβ3DETR end-to-end transformer 3D object detector (packed point format).
-
UtoniaSegmentationβUtonia linear-probing segmentation model.
-
VoteNetDetectionβVoteNet 3D object detector (packed point format).
-
VoxelMambaDetectionβVoxel Mamba: group-free state-space 3D object detector (packed point format).
-
VoxelNeXtDetectionβVoxelNeXt fully sparse 3D object detector (packed point format).
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 itsregister_modeldecorator. -
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
namefortask.
ClassificationModel [source] ΒΆ
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;
0drops the head.
Methods:
-
configure_headβBuild and return the classification head for the current
num_classes(nn.Identitywhen 0). -
reset_classifierβReplace the classification head for
num_classesoutputs (0drops 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 ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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).
DetectionModel [source] ΒΆ
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_classesoutputs. -
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
forwardoutput into packed boxes.
Attributes:
-
num_features(int) βChannel count \(C\) of the
forward_featuresoutput.
num_features abstractmethod property ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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).
PartSegmentationModel [source] ΒΆ
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;
0drops the head. -
num_categories(int) βNumber of object categories of the one-hot
categoryinput.
Methods:
-
configure_headβBuild and return the segmentation head for the current
num_classes(nn.Identitywhen 0). -
reset_classifierβReplace the segmentation head for
num_classesoutputs (0drops it). -
forwardβRun the model on a packed point cloud.
-
forward_featuresβEncode a packed point cloud, keeping what
forward_decoderneeds 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 ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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.
PretrainingModel [source] ΒΆ
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] ΒΆ
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;
0drops the head.
Methods:
-
configure_headβBuild and return the segmentation head for the current
num_classes(nn.Identitywhen 0). -
reset_classifierβReplace the segmentation head for
num_classesoutputs (0drops it). -
forwardβRun the model on a packed point cloud.
-
forward_featuresβEncode a packed point cloud, keeping what
forward_decoderneeds 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 ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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] ΒΆ
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.
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
Nonewhen 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 forhttps://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 (seeload_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 fromurlbefore 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:
-
configure_encoderβBuild the
PointTransformerV3Encoderbackbone.
Attributes:
-
num_features(int) βChannel count \(C\) entering the head: every encoder stage unpooled and concatenated.
num_features property ΒΆ
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
conv5layer 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
Nonewhen its channels are unset. -
configure_encoderβBuild the
DGCNNEncoderbackbone. -
configure_projβBuild the projection MLP applied before pooling, or
Nonewhenproj_channelsis unset.
Attributes:
-
num_features(int) βChannel count \(C\) of the pooled features entering the head.
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:
-
forward_decoderβDecode encoder features back to per-point resolution.
-
configure_stnetβBuild the dynamic spatial transformer, or
Nonewhenstnet_edge_channelsis unset. -
configure_encoderβBuild the
DGCNNEncoderbackbone. -
configure_projβBuild the projection MLP producing the global feature.
-
configure_cat_embedβBuild the MLP embedding the one-hot shape category.
Attributes:
-
num_features(int) βChannel count \(C\) entering the head: encoder features, global feature and category embedding.
num_features property ΒΆ
Channel count \(C\) entering the head: encoder features, global feature and category embedding.
forward_decoder [source] ΒΆ
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.
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
Nonewhen its channels are unset. -
configure_encoderβBuild the
DGCNNEncoderbackbone. -
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 ΒΆ
Channel count \(C\) entering the head: encoder features plus the broadcast global feature.
forward_decoder [source] ΒΆ
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.
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, orNonewhenstem_channelsis unset. -
configure_encoderβBuild the
KPFCNNEncoderbackbone. -
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.
configure_stem [source] ΒΆ
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: bool = False,
) -> Any
Forward pass of the encoder, returning pre-pooling features.
Parameters:
-
x(OptTensor) βPoint features of shape \((N, C)\). If
None,posis used as features, which requiresin_channelsto match the dimension ofpos. -
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, wherexhas shape \((N', \text{encoder\_channels}[-1])\) and \(N'\) is the number of downsampled points. Ifreturn_intermediates=True, the per-stage intermediates are appended to the tuple.
forward_head [source] ΒΆ
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 pass of the classification model.
Parameters:
-
x(OptTensor) βPoint features of shape \((N, C)\). If
None,posis used as features, which requiresin_channelsto match the dimension ofpos. -
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, orNonewhenstem_channelsis unset. -
configure_encoderβBuild the
KPFCNNEncoderbackbone. -
configure_decoderβBuild the
PointNet2Decoderupsampling the coarsest features back through the encoder skips.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_stem [source] ΒΆ
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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the BEV feature map entering the head.
reset_classifier [source] ΒΆ
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(out: TransFusionHeadOutput) -> Detection3D
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
OctreePatchEmbedstem. -
configure_encoderβBuild the
OctFormerEncoderbackbone. -
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.
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:
-
configure_stemβBuild the
OctreePatchEmbedstem. -
configure_encoderβBuild the
OctFormerEncoderbackbone. -
configure_decoderβBuild the
OctFormerDecoder, with one upsampling per stem downsampling.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the features entering the head.
num_features property ΒΆ
Feature dimension \(C\) of the features entering the head.
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
Methods:
-
configure_encoderβBuild the
PointBERTEncoderbackbone.
Attributes:
-
num_features(int) βChannel count \(C\) of the pooled features entering the head.
num_features property ΒΆ
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
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] ΒΆ
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})\) andlogits\((B, G, \text{num\_tokens})\).
Example
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] ΒΆ
Build the linear bridge from the token-embedding dimension to the transformer dimension.
configure_pos_embed [source] ΒΆ
Build the positional-embedding MLP mapping patch centers to the transformer dimension.
configure_blocks [source] ΒΆ
Build the transformer blocks with a linearly scaled stochastic-depth schedule.
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:
-
configure_h_encoderβBuild the hierarchical transformer encoder.
Attributes:
-
num_features(int) βChannel count \(C\) of the pooled features entering the head.
num_features property ΒΆ
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 ΒΆ
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.
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] ΒΆ
Build the per-stage decoder transformer blocks with a linearly scaled stochastic-depth schedule.
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 ΒΆ
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] ΒΆ
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 ΒΆ
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] ΒΆ
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] ΒΆ
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βBuild the masked transformer encoding the visible group tokens.
-
configure_decoder_pos_embedβBuild the decoder positional-embedding MLP mapping group centers to token channels.
-
configure_MAE_decoderβBuild the transformer decoder reconstructing the masked tokens.
-
configure_increase_dimβBuild the per-token linear head predicting the coordinates of each masked group.
configure_MAE_encoder [source] ΒΆ
configure_MAE_encoder() -> MaskTransformer
Build the masked transformer encoding the visible group tokens.
configure_decoder_pos_embed [source] ΒΆ
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.
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:
-
configure_encoderβBuild the
PointMambaEncoderbackbone.
Attributes:
-
num_features(int) βChannel count \(C\) of the pooled features entering the head.
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
PointMambaEncoderMAEbackbone, which encodes only the visible patches. -
configure_decoderβBuild the
PointMambaDecoderMAEthat 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
PointTransformerBlockblocks 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
Noneto fall back topos. - pos: \((N, 3)\) point coordinates.
- batch: \((N,)\) per-point batch index.
- output: \((B, \text{num\_classes})\) class logits.
Methods:
-
configure_embeddingsβBuild the embedding MLP lifting the input features to the first encoder width.
-
configure_encoderβBuild the
PointTransformerEncoderbackbone.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the encoder output.
configure_embeddings [source] ΒΆ
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
PointTransformerBlockblocks 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
PointTransformerBlockblocks 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
Noneto fall back topos. - 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
PointTransformerEncoderbackbone. -
configure_decoderβBuild the
PointTransformerDecoderupsampling the coarsest features back through the encoder skips.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_embeddings [source] ΒΆ
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
PointTransformerV2Encoderbackbone. -
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.
configure_stem [source] ΒΆ
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: 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_intermediatesisTrue, a list of dictionaries containing the intermediate features, coordinates, batch indices and pooling inverse for each encoder block.
forward_head [source] ΒΆ
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 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
PointTransformerV2Encoderbackbone. -
configure_decoderβBuild the
PointTransformerV2Decoderupsampling 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 ΒΆ
Channel count \(C\) of the per-point decoder features entering the head.
configure_stem [source] ΒΆ
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] ΒΆ
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: 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_intermediatesisTrue, a list of dictionaries containing the intermediate features, coordinates, batch indices and pooling inverse for each encoder block.
forward_head [source] ΒΆ
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 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-valuedposargument 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
PointTransformerV3Encoderbackbone. -
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.
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").
Nonekeeps the current pooling. -
**kwargs(Any, default:{}) βAdditional keyword arguments to pass to the classification head.
forward_head [source] ΒΆ
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-valuedposargument 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 requiresflash-attnand a CUDA device; passuse_flash_attn=Falseto run without it (the xCPE sparse convolution still needs aspconvbuild 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
PointTransformerV3Encoderbackbone. -
configure_decoderβBuild the
PointTransformerV3Decoderupsampling 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 ΒΆ
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] ΒΆ
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:
-
configure_encoderβBuild the
PointCNNEncoderbackbone.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the features entering the head.
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:
-
configure_encoderβBuild the
PointCNNEncoderbackbone. -
configure_decoderβBuild the
PointCNNDecoder, mirroring the decimating encoder stages in reverse.
Attributes:
-
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:
-
configure_encoderβBuild the
PointConvDensityEncoderbackbone.
Attributes:
-
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
Methods:
-
configure_encoderβBuild the mini-PointNet patch embedder tokenizing each patch.
-
configure_pos_embedβBuild the sinusoidal positional embedding of the patch centers.
-
configure_blocksβBuild the causally-masked GPT extractor.
-
configure_generator_blocksβBuild the GPT generator predicting the next patch (unused at finetuning, kept for weight compatibility).
Attributes:
-
num_features(int) βChannel count \(C\) of the pooled features entering the head.
num_features property ΒΆ
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βBuild the mini-PointNet patch embedder tokenizing each patch.
-
configure_pos_embedβBuild the sinusoidal positional embedding of the patch centers and relative directions.
-
configure_blocksβBuild the causally-masked GPT extractor.
-
configure_generator_blocksβBuild the GPT generator predicting the next patch from the extractor features.
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
PointMLPEncoderbackbone.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the encoder output.
configure_stem [source] ΒΆ
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
PointMLPEncoderbackbone. -
configure_decoderβBuild the
PointMLPDecoder, mirroring the encoder channels in reverse.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_stem [source] ΒΆ
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
PointNetEncoderbackbone. -
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.
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");
Nonekeeps the current pooling. -
**kwargs(Any, default:{}) βAdditional keyword arguments to pass to the classification head.
forward_features [source] ΒΆ
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 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 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
PointNetEncoderbackbone. -
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 ΒΆ
Channel count \(C\) entering the head: per-point features concatenated with the global feature.
forward_decoder [source] ΒΆ
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");
Nonekeeps the current pooling. -
**kwargs(Any, default:{}) βAdditional keyword arguments to pass to the segmentation head.
forward_features [source] ΒΆ
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)wherepre_pooling_featuresis of shape \((N, \text{mlp2\_dims}[-1])\) andpoint_featuresis 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 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
PointNet2Encoderbackbone. -
configure_aggrβBuild the aggregation MLP applied to the encoder output, or
Nonewhenaggr_channelsis unset.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the encoder output, after the optional aggregation MLP.
num_features property ΒΆ
Feature dimension \(C\) of the encoder output, after the optional aggregation MLP.
configure_encoder [source] ΒΆ
configure_encoder() -> PointNet2Encoder
Build the PointNet2Encoder backbone.
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
PointNet2Encoderbackbone. -
configure_aggrβBuild the aggregation MLP applied to the encoder output, or
Nonewhenaggr_channelsis unset. -
configure_decoderβBuild the
PointNet2Decoderupsampling the coarsest features back through the encoder skips.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_encoder [source] ΒΆ
configure_encoder() -> PointNet2Encoder
Build the PointNet2Encoder backbone.
configure_aggr [source] ΒΆ
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
Nonewhenstem_channelsis unset. -
configure_encoderβBuild the
PointNeXtEncoderbackbone. -
configure_global_saβBuild the shared MLP over
cat([pos, x])applied before pooling, orNonewhenglobal_sa_channelsis unset.
Attributes:
-
num_features(int) βFeature dimension \(C\) fed to the classification head, after the optional global set-abstraction.
num_features property ΒΆ
Feature dimension \(C\) fed to the classification head, after the optional global set-abstraction.
configure_stem [source] ΒΆ
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.
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
Nonewhenstem_channelsis unset. -
configure_encoderβBuild the
PointNeXtEncoderbackbone. -
configure_decoderβBuild the
PointNeXtPartDecoderupsampling 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 ΒΆ
Channel count \(C\) entering the head: decoder features concatenated with their global max and mean pools.
configure_stem [source] ΒΆ
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
Nonewhenstem_channelsis unset. -
configure_encoderβBuild the
PointNeXtEncoderbackbone. -
configure_decoderβBuild the
PointNeXtDecoderupsampling the coarsest features back through the encoder skips.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_stem [source] ΒΆ
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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the BEV feature map entering the head.
reset_classifier [source] ΒΆ
Replace the classification branch of the detection head for num_classes outputs.
Models whose head is not rebuildable in isolation raise NotImplementedError.
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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the BEV feature map entering the head.
reset_classifier [source] ΒΆ
Replace the classification branch of the detection head for num_classes outputs.
Models whose head is not rebuildable in isolation raise NotImplementedError.
configure_head [source] ΒΆ
configure_head() -> MultiGroupAnchorHead
Build the multi-group anchor head.
decode [source] ΒΆ
decode(out: AnchorHeadMultiOutput) -> Detection3D
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 (
PointNet2Decoderorder). -
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 (
-1groups 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_threshto 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 ΒΆ
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:
-
Union[PointRCNNTrainOutput, PointRCNNOutput]βA
PointRCNNTrainOutputin train mode (stage-1 point predictions plus sampled-ROI stage-2 -
Union[PointRCNNTrainOutput, PointRCNNOutput]βtensors for the loss), otherwise a
PointRCNNOutput(refined boxes with confidences).
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.
Noneskips 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
Noneuses 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
Nonewhenglobal_channelsis empty.
Attributes:
-
num_features(int) βChannel count \(C\) of the pooled features entering the head.
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_dimwhen 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.
Noneskips 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
Noneuses 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.
Noneuses 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
Nonewhenglobal_channelsis empty.
Attributes:
-
num_features(int) βChannel count \(C\) entering the head: every layer output concatenated with the global feature.
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
Noneuses 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; whenNone,posis used as features.pos: \((N, 3)\) point coordinates.batch: \((N,)\) batch indices.- output: \((B, \text{num\_classes})\) classification logits.
Methods:
-
configure_encoderβBuild the
PVCNN2Encoderbackbone.
Attributes:
-
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
Noneuses 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
Noneuses 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_channelsis set: the MLP head applies its ownhead_dropout. -
head_channels(Optional[Sequence[int]], default:None) βHidden widths of the segmentation head MLP.
Noneuses a single linear layer. -
head_dropout(float, default:0.0) βDropout probability inside the head MLP.
Shape
x: \((N, C_{in})\) point features; whenNone,posis 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
PVCNN2Encoderbackbone. -
configure_decoderβBuild the
PVCNN2Decoderupsampling the coarsest features back through the encoder skips.
Attributes:
-
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
Noneto 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, orNonewhenstem_channelsis unset. -
configure_encoderβBuild the
RandLANetEncoderbackbone. -
configure_aggrβBuild the aggregation MLP applied to the encoder output, or
Nonewhenaggr_channelsis unset.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the encoder output, after the optional aggregation MLP.
num_features property ΒΆ
Feature dimension \(C\) of the encoder output, after the optional aggregation MLP.
configure_stem [source] ΒΆ
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.
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
Noneto 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 aCallable/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 MLPact_firstsemantics):Linear β Act β Norm β Dropoutinstead of the defaultLinear β 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=Truesince it has no normalization.
Methods:
-
configure_stemβBuild the stem lifting the input features to
stem_channels, orNonewhenstem_channelsis unset. -
configure_encoderβBuild the
RandLANetEncoderbackbone. -
configure_aggrβBuild the aggregation MLP applied to the encoder output, or
Nonewhenaggr_channelsis unset. -
configure_decoderβBuild the
RandLANetDecoderupsampling the bottleneck features back through the encoder skips.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_stem [source] ΒΆ
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] ΒΆ
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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the BEV feature map entering the head.
reset_classifier [source] ΒΆ
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.
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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the BEV feature map entering the head.
reset_classifier [source] ΒΆ
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_head [source] ΒΆ
configure_head() -> MultiGroupAnchorHead
Build the multi-group anchor head.
decode [source] ΒΆ
decode(out: AnchorHeadMultiOutput) -> Detection3D
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:
-
configure_encoderβBuild the
PointTransformerV3Encoderbackbone. -
forwardβForward pass.
Attributes:
-
num_features(int) βChannel count \(C\) entering the head: every encoder stage unpooled and concatenated.
num_features property ΒΆ
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 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;
0yields 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
intis 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 ΒΆ
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:
-
configure_input_convβBuild the submanifold stem convolution.
-
configure_unetβBuild the recursive sparse UNet with windowed attention at every level.
-
configure_output_layerβBuild the final normalization and activation applied before the head.
Attributes:
-
num_features(int) βChannel count \(C\) of the full-resolution features entering the head.
num_features property ΒΆ
Channel count \(C\) of the full-resolution features entering the head.
configure_unet [source] ΒΆ
configure_unet() -> SphereFormerUBlock
Build the recursive sparse UNet with windowed attention at every level.
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
xisNone. -
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 ΒΆ
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_dimwhenNone. -
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
SPVCNNEncoderbackbone.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the encoder output.
configure_stem [source] ΒΆ
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_dimwhenNone. -
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
SPVCNNEncoderbackbone. -
configure_decoderβBuild the
SPVCNNDecoderupsampling the coarsest features back through the encoder skips.
Attributes:
-
num_features(int) βFeature dimension \(C\) of the decoder output.
configure_stem [source] ΒΆ
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 keepspreencoder_num_pointstokens) 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:
-
configure_pre_encoderβBuild the set-abstraction tokenizer.
-
configure_encoderβBuild the transformer encoder (masked with one interim downsampling for 3DETR-m).
-
configure_encoder_to_decoder_projectionβBuild the projection from encoder tokens to the decoder dimension.
-
configure_pos_embeddingβBuild the Fourier position embedding.
-
configure_query_projectionβBuild the query projection.
-
configure_decoderβBuild the transformer decoder.
-
configure_mlp_headsβBuild the per-query class, center, size and heading heads.
-
run_encoderβTokenizes the point cloud with the pre-encoder and runs the masked transformer encoder.
-
get_query_embeddingsβSamples the query positions among the encoder tokens and embeds them into decoder queries.
-
decodeβDecode a forward output into raw per-query detections (no NMS, threshold, or filtering).
Attributes:
-
num_features(int) βChannel count \(C\) of the encoder tokens entering the decoder.
num_features property ΒΆ
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] ΒΆ
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.
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
Noneto 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:
-
out(ThreeDETROutput) βA
ThreeDETROutputfromforward.
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_probsholds -
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
PointTransformerV3Encoderbackbone 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 ΒΆ
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 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:
-
configure_backboneβBuild the PointNet++ seed backbone.
-
configure_vgenβBuild the voting module.
-
configure_proposalβBuild the proposal module.
-
decodeβDecode a forward output into raw per-proposal detections (no NMS, threshold, or filtering).
Attributes:
-
num_features(int) βChannel count \(C\) of the per-seed backbone features entering the voting head.
num_features property ΒΆ
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_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:
-
out(VoteNetOutput) βA
VoteNetOutputfromforward.
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_probsholds -
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
RMSNorminside 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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the BEV feature map entering the head.
reset_classifier [source] ΒΆ
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.
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
CenterHeadOutputfromforward. -
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_classesoutputs. -
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 ΒΆ
Channel count \(C\) of the sparse voxel features entering the head.
reset_classifier [source] ΒΆ
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: 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-segmentationordetection). Optional whennameis 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.savefiles,.safetensors, and Lightning checkpoints (the wrapped network is extracted). -
return_info(bool, default:False) βAlso return the registry entry (
ModelInfoDict), with the resolvedtaskand the effectivehparams. -
**kwargs(Any, default:{}) βOverrides merged into the registered
hparamsand passed to the model constructor.
Returns:
-
AnyβThe model, or a
(model, info)tuple whenreturn_infois true.
Raises:
-
ValueErrorβIf
taskornameis unknown, iftaskis omitted for a name registered under several tasks, if bothpretrainedandcheckpoint_pathare passed, or ifpretrainedis requested for an entry that registers no weights. -
TypeErrorβIf the entry registers architecture hparams only and the data-dependent arguments (typically
in_channelsandnum_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:
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;
Noneleaves 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_inorremoved_inis not a version string, ifremoved_inis not a later release thandeprecated_in, if the decorated callable is not registered under exactly one name (the decorator belongs aboveregister_model), or ifnameis 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 (
fnmatchsyntax, e.g."pointnext*"). -
task(Optional[Task], default:None) βRestrict the listing to one registry;
Nonelists 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:
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
nameis registered under a single task.
Returns:
-
ModelInfoDictβThe entry as a
ModelInfoDict, with the registeredhparams.
Raises:
-
ValueErrorβIf
taskornameis unknown, or iftaskis omitted for a name registered under several tasks.
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-segmentationordetection). -
hparams(Optional[Dict[str, Any]], default:None) βDefault keyword arguments the entry point is called with;
create_modelkwargs 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
forwardtakes positionally, in order; a dotted key reads an attribute (e.g.octree.depth).utils.data.select_inputsgathers them from a collated batch. -
weights(Union[str, WeightsDict, None], default:None) βPretrained checkpoint, either a URL string or a
WeightsDictwith metadata.
Returns:
-
CallableβThe decorator registering its target callable.
Raises:
-
ValueErrorβIf
taskis unknown, or ifnameis the former name of another model oftask.
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)