pointnet

PointNet classification and segmentation models.

First page of PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation

1612.00593 · December 2016

Classes:

PointNetEncoder –

PointNet encoder module that processes point clouds to extract global feature vectors as described in the original PointNet paper
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

PointNetEncoder [source] ¶

PointNetEncoder(
    spatial_dim: int = 3,
    in_channels: int = 0,
    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,
)

Bases: Module

PointNet encoder module that processes point clouds to extract global feature vectors as described in the original PointNet paper PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation.

The encoder follows the PointNet architecture by:

Applying a spatial transformer network (T-Net) to align input point coordinates,
Processing points through the first MLP to extract per-point features,
Optionally applying a feature transformer network to align feature space,
Processing through the second MLP to extract higher-level features.

Abstract

This is the core feature extraction component of PointNet. The global features can be used for classification tasks, while the combination of global and point features can be used for segmentation tasks.

Tip

To get actual global features, you should apply your own pooling operation on the output of this module, like:

import torch
from torch_geometric.utils import scatter

from torch_pointcloud.models.pointnet import PointNetEncoder

encoder = PointNetEncoder()
pos = torch.rand(2048, 3)
batch = torch.arange(2).repeat_interleave(1024)
out = encoder(None, pos, batch)
global_features = scatter(out, batch, dim=0, reduce="max")

Parameters:

spatial_dim (int, default: 3 ) –

Dimension of point coordinates.
in_channels (int, default: 0 ) –

Dimension of additional point features.
mlp1_channels (Sequence[int], default: (64,) ) –

Dimensions of the first MLP.
mlp2_channels (Sequence[int], default: (128, 1024) ) –

Dimensions of the 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 the feature transformer network.
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.

Methods:

forward –

Forward pass of the PointNet encoder.

forward [source] ¶

forward(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

forward(
    x: Optional[Tensor],
    pos: Tensor,
    batch: Tensor,
    return_point_features: bool = True,
) -> Tuple[Tensor, Tensor]

forward(
    x: Optional[Tensor],
    pos: Tensor,
    batch: Tensor,
    return_point_features: bool = False,
) -> Union[Tensor, Tuple[Tensor, Tensor]]

Forward pass of the PointNet encoder.

Parameters:

x (Optional[Tensor]) –

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

Point coordinates of shape \((N, D)\).
batch (Tensor) –

Batch indices for each point of shape \((N,)\).
return_point_features (bool, default: False ) –

Whether to return per-point features.

Returns:

Tensor –

If return_point_features=False, per-point features of shape \((N, C_2)\) where \(N\) is the number of points and \(C_2\) is the last dimension of the second MLP.
Tuple[Tensor, Tensor] –
If return_point_features=True, a tuple of:
- x of shape \((N, C_2)\) where \(C_2\) is the last dimension of the second MLP.
- point_features of shape \((N, C_1)\) where \(C_1\) is the last dimension of the first MLP.

PointNetClassification [source] ¶

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

Bases: ClassificationModel

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

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

Abstract

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

Tip

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

Tip

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

Parameters:

in_channels (int) –

Dimension of additional point features.
num_classes (int) –

Number of output classes.
spatial_dim (int, default: 3 ) –

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

Dropout rate applied within the classification head.
global_pool (PoolLike, default: 'max' ) –

Pooling method to aggregate point features ("max" or "mean").
mlp1_channels (Sequence[int], default: (64,) ) –

Dimensions of encoder's first MLP.
mlp2_channels (Sequence[int], default: (128, 1024) ) –

Dimensions of encoder's second MLP.
act (Union[str, Callable, None], default: 'relu' ) –

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

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

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

Keyword arguments for the normalization layers.
use_features_transform (bool, default: True ) –

Whether to use feature transformation.
tnet_mlp1_dims (Sequence[int], default: (64, 128, 1024) ) –

Dimensions of T-Net first MLP.
tnet_mlp2_dims (Sequence[int], default: (512, 256) ) –

Dimensions of T-Net second MLP.
tnet_act (Union[str, Callable, None], default: 'relu' ) –

Activation function for T-Net.
tnet_act_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the T-Net activation function.
tnet_norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization for T-Net.
tnet_norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the T-Net normalization layers.
head_channels (Sequence[int], default: (512, 256) ) –

Hidden dimensions of the classification head MLP.

Shape

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

Methods:

configure_encoder –

Build the PointNetEncoder backbone.
reset_classifier –

Resets the classification head with new parameters.
forward_features –

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

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

Forward pass of the PointNet classification network.

Attributes:

num_features (int) –

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

num_features `property` ¶

num_features: int

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

configure_encoder [source] ¶

configure_encoder() -> PointNetEncoder

Build the PointNetEncoder backbone.

reset_classifier [source] ¶

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

Resets the classification head with new parameters.

Parameters:

num_classes (int) –

Number of output classes.
global_pool (Optional[PoolLike], default: None ) –

Pooling method to aggregate point features ("max" or "mean"); None keeps the current pooling.
**kwargs (Any, default: {} ) –

Additional keyword arguments to pass to the classification head.

forward_features [source] ¶

forward_features(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

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

Parameters:

x (Optional[Tensor]) –

Additional point features of shape \((N, \text{in\_channels})\).
pos (Tensor) –

Point coordinates of shape \((N, \text{spatial\_dim})\).
batch (Tensor) –

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

Returns:

Tensor –

Pre-pooling features of shape \((N, \text{mlp2\_dims}[-1])\) where \(N\) is the number of points.

forward_head [source] ¶

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

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

Parameters:

x (Tensor) –

Pre-pooling features of shape \((N, \text{mlp2\_dims}[-1])\) where \(N\) is the number of points.
batch (Tensor) –

Batch indices for each point of shape \((N,)\).
pre_logits (bool, default: False ) –

Whether to return pre-logits. Defaults to False.

Returns:

Tensor –

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

forward [source] ¶

forward(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

Forward pass of the PointNet classification network.

Parameters:

x (Optional[Tensor]) –

Additional point features of shape \((N, \text{in\_channels})\).
pos (Tensor) –

Point coordinates of shape \((N, \text{spatial\_dim})\).
batch (Tensor) –

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

Returns:

Tensor –

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

PointNetSegmentation [source] ¶

PointNetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    dropout: float = 0.3,
    mlp1_channels: Sequence[int] = (64,),
    mlp2_channels: Sequence[int] = (128, 1024),
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    global_pool: PoolLike = "max",
    use_features_transform: bool = True,
    tnet_mlp1_dims: Sequence[int] = (64, 128, 1024),
    tnet_mlp2_dims: Sequence[int] = (512, 256),
    tnet_act: Union[str, Callable, None] = "relu",
    tnet_act_kwargs: Optional[Dict[str, Any]] = None,
    tnet_norm: Union[str, Callable, None] = "batch_norm",
    tnet_norm_kwargs: Optional[Dict[str, Any]] = None,
    seg_head_dims: Sequence[int] = (512, 256, 128),
)

Bases: SemanticSegmentationModel

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

This model implements the segmentation variant of PointNet as described in the original paper. It extracts both local point features and global shape features, combines them for each point, and predicts per-point semantic labels. This architecture enables the network to consider both local geometry and global context for segmentation.

Abstract

The key innovation in the segmentation variant is the concatenation of global features with per-point features, allowing each point's classification to be informed by both local geometry and the global shape context. This enables part segmentation that is aware of the overall object structure.

Parameters:

in_channels (int) –

Dimension of additional point features.
num_classes (int) –

Number of segmentation classes.
spatial_dim (int, default: 3 ) –

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

Dropout rate applied within the segmentation head.
mlp1_channels (Sequence[int], default: (64,) ) –

Dimensions of encoder's first MLP.
mlp2_channels (Sequence[int], default: (128, 1024) ) –

Dimensions of encoder's second MLP.
act (Union[str, Callable, None], default: 'relu' ) –

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

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

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

Keyword arguments for the normalization layers.
global_pool (PoolLike, default: 'max' ) –

Pooling method for global features ("max" or "mean").
use_features_transform (bool, default: True ) –

Whether to use feature transformation.
tnet_mlp1_dims (Sequence[int], default: (64, 128, 1024) ) –

Dimensions of T-Net first MLP.
tnet_mlp2_dims (Sequence[int], default: (512, 256) ) –

Dimensions of T-Net second MLP.
tnet_act (Union[str, Callable, None], default: 'relu' ) –

Activation function for T-Net.
tnet_act_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the T-Net activation function.
tnet_norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization for T-Net.
tnet_norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the T-Net normalization layers.
seg_head_dims (Sequence[int], default: (512, 256, 128) ) –

Dimensions of segmentation head MLPs.

Shape

Input: features of shape \((N, \text{in\_channels})\) (optional), points of shape \((N, \text{spatial\_dim})\)
Output: logits of shape \((N, \text{num\_classes})\)

Methods:

configure_encoder –

Build the PointNetEncoder backbone.
reset_classifier –

Resets the segmentation head with new parameters.
forward_features –

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

Forward pass of the segmentation head from pre-pooling features.
forward –

Forward pass of the PointNet segmentation network.
forward_decoder –

Decode encoder features back to per-point resolution.

Attributes:

num_features (int) –

Channel count \(C\) entering the head: per-point features concatenated with the global feature.

num_features `property` ¶

num_features: int

Channel count \(C\) entering the head: per-point features concatenated with the global feature.

configure_encoder [source] ¶

configure_encoder() -> PointNetEncoder

Build the PointNetEncoder backbone.

reset_classifier [source] ¶

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

Resets the segmentation head with new parameters.

Parameters:

num_classes (int) –

Number of output classes.
global_pool (Optional[PoolLike], default: None ) –

Pooling method for global features ("max" or "mean"); None keeps the current pooling.
**kwargs (Any, default: {} ) –

Additional keyword arguments to pass to the segmentation head.

forward_features [source] ¶

forward_features(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tuple[Tensor, Tensor]

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

Parameters:

x (Optional[Tensor]) –

Additional point features of shape \((N, \text{in\_channels})\).
pos (Tensor) –

Point coordinates of shape \((N, \text{spatial\_dim})\).
batch (Tensor) –

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

Returns:

Tuple[Tensor, Tensor] –

A tuple of (pre_pooling_features, point_features) where pre_pooling_features is of shape \((N, \text{mlp2\_dims}[-1])\) and point_features is of shape \((N, \text{mlp1\_dims}[-1])\).

forward_head [source] ¶

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

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

Parameters:

x (Tensor) –

Pre-pooling features of shape \((N, \text{mlp2\_dims}[-1])\) where \(N\) is the number of points.
point_features (Tensor) –

Point features of shape \((N, \text{mlp1\_dims}[-1])\).
batch (Tensor) –

Batch indices for each point of shape \((N,)\).
pre_logits (bool, default: False ) –

Whether to return pre-logits. Defaults to False.

Returns:

Tensor –

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

forward [source] ¶

forward(
    x: Optional[Tensor], pos: Tensor, batch: Tensor
) -> Tensor

Forward pass of the PointNet segmentation network.

Parameters:

x (Optional[Tensor]) –

Additional point features of shape \((N, \text{in\_channels})\).
pos (Tensor) –

Point coordinates of shape \((N, \text{spatial\_dim})\).
batch (Tensor) –

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

Returns:

Tensor –

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

forward_decoder [source] ¶

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

Decode encoder features back to per-point resolution.

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

pointnet

PointNetEncoder [source] ¶

forward [source] ¶

PointNetClassification [source] ¶

num_features property ¶

configure_encoder [source] ¶

reset_classifier [source] ¶

forward_features [source] ¶

forward_head [source] ¶

forward [source] ¶

PointNetSegmentation [source] ¶

num_features property ¶

configure_encoder [source] ¶

reset_classifier [source] ¶

forward_features [source] ¶

forward_head [source] ¶

forward [source] ¶

forward_decoder [source] ¶

num_features `property` ¶

num_features `property` ¶