Skip to content

Point-Mamba

Point-Mamba classification and masked autoencoder pretraining models.

First page of PointMamba: A Simple State Space Model for Point Cloud Analysis

2402.10739 · February 2024

Classes:

Functions:

  • order_sort –

    Sort points along a space-filling curve, so a state space model scans spatially close points in sequence.

PointMambaEncoderMAE 

PointMambaEncoderMAE(
    in_channels: int,
    embed_dim: int,
    depth: int,
    num_group: int,
    group_size: int,
    mask_ratio: float,
    drop_path: float = 0.0,
    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: Module

PointMamba encoder for masked-autoencoder pre-training.

Embeds patches, serializes them along a randomly chosen Hilbert or Trans-Hilbert curve, masks a fraction of the patch tokens, and processes only the visible tokens with the Mamba blocks. The forward pass returns a dict with the visible-token features, the densified patch positions and reconstruction targets, and the visible / masked token indices.

Parameters:

  • in_channels (int) –

    The number of input channels.

  • embed_dim (int) –

    The number of token channels.

  • depth (int) –

    The number of Mamba blocks.

  • num_group (int) –

    The number of patches to sample.

  • group_size (int) –

    The number of neighbors to consider for each patch.

  • mask_ratio (float) –

    The fraction of patch tokens to mask.

  • drop_path (float, default: 0.0 ) –

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

  • 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: a dict with x_vis \((B, V, C)\), pos_dense \((B, P, 3)\), target_pos_dense \((B, P, M, 3)\), mask_idx \((B, P - V)\), vis_idx \((B, V)\), and the number of patches \(P\).

PointMambaDecoderMAE 

PointMambaDecoderMAE(
    embed_dim: int,
    depth: int,
    drop_path: float,
    spatial_dim: int = 3,
    pos_embed_channels: Sequence[int] = (128,),
    bias: Union[bool, List[bool]] = True,
)

Bases: Module

PointMamba decoder for masked-autoencoder pre-training.

Scatters the visible-token features into a full-length sequence of learnable mask tokens, adds the positional embedding of every patch center, and refines the sequence with a stack of Mamba blocks.

Parameters:

  • embed_dim (int) –

    The number of token channels.

  • depth (int) –

    The number of Mamba blocks.

  • drop_path (float) –

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

  • spatial_dim (int, default: 3 ) –

    The dimension of the spatial features.

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

    Hidden widths of the positional-embedding MLP.

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

    Whether to use bias in the MLPs.

Shape
  • Input: x_vis \((B, V, C)\), pos_dense \((B, P, 3)\), and ids_keep \((B, V)\).
  • Output: \((B, P, C)\).

PointMambaClassification 

PointMambaClassification(
    in_channels: int,
    num_classes: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_group: int = 64,
    group_size: int = 32,
    drop_path: float = 0.1,
    use_cls_token: bool = False,
    spatial_dim: int = 3,
    patch_local_channels: Sequence[int] = (128, 256),
    patch_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: Union[bool, List[bool]] = True,
    dropout: float = 0.5,
    global_pool: AdaptivePoolLike = "mean",
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
)

Bases: ClassificationModel

Methods:

Attributes:

  • num_features (int) –

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

num_features property 

num_features: int

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

configure_encoder 

configure_encoder() -> PointMambaEncoder

Build the PointMambaEncoder backbone.

PointMambaMAE 

PointMambaMAE(
    in_channels: int,
    *,
    embed_dim: int = 384,
    encoder_depth: int = 12,
    decoder_depth: int = 4,
    num_group: 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: BaseModel

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 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_group (int, default: 64 ) –

    The number of patches to sample.

  • group_size (int, default: 32 ) –

    The number of neighbors to consider for each patch.

  • mask_ratio (float, default: 0.6 ) –

    The fraction of patch tokens to mask.

  • drop_path (float, default: 0.1 ) –

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

  • spatial_dim (int, default: 3 ) –

    The dimension of the spatial features.

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

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

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

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

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

    Hidden widths of the positional-embedding MLPs.

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

    The activation function to use.

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

    The keyword arguments to pass to the activation function.

  • act_first (bool, default: False ) –

    Whether to apply the activation function before the normalization.

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

    The normalization function to use.

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

    The keyword arguments to pass to the normalization function.

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

    Whether to use bias in the MLPs.

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

Methods:

  • configure_encoder –

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

  • configure_decoder –

    Build the PointMambaDecoderMAE that reconstructs the masked patches.

configure_encoder 

configure_encoder() -> PointMambaEncoderMAE

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

configure_decoder 

configure_decoder() -> PointMambaDecoderMAE

Build the PointMambaDecoderMAE that reconstructs the masked patches.

order_sort 

order_sort(
    pos_grid: Tensor,
    batch: Tensor,
    order: SerializationOrder,
) -> Tensor

Sort points along a space-filling curve, so a state space model scans spatially close points in sequence.

Parameters:

  • pos_grid (Tensor) –

    Non-negative integer grid coordinates of shape \((N, 3)\).

  • batch (Tensor) –

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

  • order (SerializationOrder) –

    Space-filling curve to encode the coordinates along.

Returns:

  • Tensor –

    The permutation of shape \((N,)\) that sorts the points by their serialization code.

Raises:

  • ValueError –

    If pos_grid holds a negative coordinate, which would silently wrap around to a valid code.