point_mae

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

First page of Masked Autoencoders for Point Cloud Self-supervised Learning

2203.06604 · March 2022

Classes:

TransformerEncoder –

Plain (non-hierarchical) transformer encoder, as in [Masked Autoencoders for Point Cloud
TransformerDecoder –

Transformer decoder, as in [Masked Autoencoders for Point Cloud Self-supervised
MaskTransformer –

Masked patch-embedding transformer encoder, as in [Masked Autoencoders for Point Cloud
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

TransformerEncoder [source] ¶

TransformerEncoder(
    embed_dim: int = 768,
    depth: int = 4,
    num_heads: int = 12,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = False,
    dropout: float = 0.0,
    attn_dropout: float = 0.0,
    drop_path: Union[float, List[float]] = 0.0,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = LayerNorm,
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Plain (non-hierarchical) transformer encoder, as in Masked Autoencoders for Point Cloud Self-supervised Learning, adapted from Pang-Yatian/Point-MAE.

A stack of pre-norm transformer blocks. The positional embedding is added to the tokens before every block. The encoder can optionally return the hidden states at a fixed set of block indices (used by the segmentation decoder).

Parameters:

embed_dim (int, default: 768 ) –

Token channels.
depth (int, default: 4 ) –

Number of transformer blocks.
num_heads (int, default: 12 ) –

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

Hidden-channel expansion ratio of the MLP.
qkv_bias (bool, default: False ) –

Whether to use a bias term in the query / key / value projection.
dropout (float, default: 0.0 ) –

Dropout rate inside the blocks.
attn_dropout (float, default: 0.0 ) –

Dropout rate applied to the attention matrix.
drop_path (Union[float, List[float]], default: 0.0 ) –

Stochastic-depth drop-path rate(s). A float applies to every block.
act (Union[str, Callable, None], default: 'gelu' ) –

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

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

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

Extra arguments for the normalization.

Shape

Input: \((B, N, C)\) where \(B\) is the batch size, \(N\) is the sequence length, and \(C\) is embed_dim.
Output: \((B, N, C)\).

TransformerDecoder [source] ¶

TransformerDecoder(
    embed_dim: int = 384,
    depth: int = 4,
    num_heads: int = 6,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = False,
    dropout: float = 0.0,
    attn_dropout: float = 0.0,
    drop_path: Union[float, List[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,
)

Bases: Module

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

A stack of pre-norm transformer blocks followed by a final layer normalization. Only the hidden states of the last return_token_num tokens are returned (the masked tokens).

Parameters:

embed_dim (int, default: 384 ) –

Token channels.
depth (int, default: 4 ) –

Number of transformer blocks.
num_heads (int, default: 6 ) –

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

Hidden-channel expansion ratio of the MLP.
qkv_bias (bool, default: False ) –

Whether to use a bias term in the query / key / value projection.
dropout (float, default: 0.0 ) –

Dropout rate inside the blocks.
attn_dropout (float, default: 0.0 ) –

Dropout rate applied to the attention matrix.
drop_path (Union[float, List[float]], default: 0.1 ) –

Stochastic-depth drop-path rate(s). A float applies to every block.
act (Union[str, Callable, None], default: 'gelu' ) –

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

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

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

Extra arguments for the normalization.

Shape

Input: \((B, N, C)\) where \(B\) is the batch size, \(N\) is the sequence length, and \(C\) is embed_dim.
Output: \((B, R, C)\) where \(R\) is return_token_num.

MaskTransformer [source] ¶

MaskTransformer(
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    mask_ratio: float = 0.6,
    drop_path: float = 0.1,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    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,
)

Bases: Module

Masked patch-embedding transformer encoder, as in Masked Autoencoders for Point Cloud Self-supervised Learning, adapted from Pang-Yatian/Point-MAE.

Tokenizes local groups with the mini-PointNet encoder, randomly masks a fraction of the tokens, and processes only the visible tokens with a plain transformer encoder. This is the encoder of the Point-MAE pretraining model.

Parameters:

embed_dim (int, default: 384 ) –

Token channels.
depth (int, default: 12 ) –

Number of transformer blocks.
num_heads (int, default: 6 ) –

Number of attention heads.
mask_ratio (float, default: 0.6 ) –

Fraction of tokens to mask.
drop_path (float, default: 0.1 ) –

Stochastic-depth drop-path rate.
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.
spatial_dim (int, default: 3 ) –

Spatial dimension of the input point cloud.
act (Union[str, Callable, None], default: 'gelu' ) –

Activation function of the transformer MLPs and the positional embedding.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

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

Normalization applied inside the transformer blocks.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the normalization.

Shape

Input: \((B, G, M, 3)\) and \((B, G, 3)\).
Output: visible tokens \((B, V, C)\) and a boolean mask \((B, G)\).

Methods:

mask_center_rand –

Draw a random mask hiding mask_ratio of the groups, independently for every sample.

mask_center_rand [source] ¶

mask_center_rand(center: Tensor) -> Tensor

Draw a random mask hiding mask_ratio of the groups, independently for every sample.

Parameters:

center (Tensor) –

Group centers of shape \((B, G, 3)\).

Returns:

Tensor –

A boolean mask of shape \((B, G)\), True where the group is masked out.

PointMAEClassification [source] ¶

PointMAEClassification(
    in_channels: int,
    num_classes: int,
    *,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    drop_path: float = 0.1,
    dropout: float = 0.5,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = LayerNorm,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    spatial_dim: int = 3,
)

Bases: ClassificationModel

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

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

Parameters:

in_channels (int) –

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

Number of output classes.
embed_dim (int, default: 384 ) –

Token channels.
depth (int, default: 12 ) –

Number of transformer blocks.
num_heads (int, default: 6 ) –

Number of attention heads.
num_groups (int, default: 64 ) –

Number of groups (FPS centers) per sample.
group_size (int, default: 32 ) –

Number of neighbors per group.
encoder_local_channels (Sequence[int], default: (128, 256) ) –

Hidden widths of the patch embedder's per-point MLP.
encoder_global_channels (Sequence[int], default: (512,) ) –

Hidden widths of the patch embedder's per-group MLP.
pos_embed_channels (Sequence[int], default: (128,) ) –

Hidden widths of the positional-embedding MLP.
drop_path (float, default: 0.1 ) –

Stochastic-depth drop-path rate.
dropout (float, default: 0.5 ) –

Dropout rate in the classification head.
act (Union[str, Callable, None], default: 'gelu' ) –

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

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

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

Extra arguments for the normalization.
spatial_dim (int, default: 3 ) –

Spatial dimension of the input point cloud.

Shape

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

Methods:

configure_encoder –

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

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

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

Attributes:

num_features (int) –

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

num_features `property` ¶

num_features: int

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

configure_encoder [source] ¶

configure_encoder() -> PointPatchEmbed

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

configure_pos_embed [source] ¶

configure_pos_embed() -> MLP

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

configure_blocks [source] ¶

configure_blocks() -> TransformerEncoder

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

PointMAEPartSegmentation [source] ¶

PointMAEPartSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    num_categories: int = 16,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    num_groups: int = 128,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    drop_path: float = 0.1,
    dropout: float = 0.5,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = LayerNorm,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    spatial_dim: int = 3,
)

Bases: PartSegmentationModel

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

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

Parameters:

in_channels (int) –

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

Number of output part classes (across all categories).
num_categories (int, default: 16 ) –

Number of object categories for the category one-hot branch.
embed_dim (int, default: 384 ) –

Token channels.
depth (int, default: 12 ) –

Number of transformer blocks.
num_heads (int, default: 6 ) –

Number of attention heads.
num_groups (int, default: 128 ) –

Number of groups (FPS centers) per sample.
group_size (int, default: 32 ) –

Number of neighbors per group.
encoder_local_channels (Sequence[int], default: (128, 256) ) –

Hidden widths of the patch embedder's per-point MLP.
encoder_global_channels (Sequence[int], default: (512,) ) –

Hidden widths of the patch embedder's per-group MLP.
pos_embed_channels (Sequence[int], default: (128,) ) –

Hidden widths of the positional-embedding MLP.
drop_path (float, default: 0.1 ) –

Stochastic-depth drop-path rate.
dropout (float, default: 0.5 ) –

Dropout rate in the per-point head.
act (Union[str, Callable, None], default: 'gelu' ) –

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

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

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

Extra arguments for the normalization.
spatial_dim (int, default: 3 ) –

Spatial dimension of the input point cloud.

Shape

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

Methods:

configure_encoder –

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

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

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

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

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

Attributes:

num_features (int) –

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

num_features `property` ¶

num_features: int

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

configure_encoder [source] ¶

configure_encoder() -> PointPatchEmbed

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

configure_pos_embed [source] ¶

configure_pos_embed() -> MLP

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

configure_blocks [source] ¶

configure_blocks() -> TransformerEncoder

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

configure_label_conv [source] ¶

configure_label_conv() -> MLP

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

configure_propagation_0 [source] ¶

configure_propagation_0() -> PointNet2FeaturePropagation

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

PointMAEPretraining [source] ¶

PointMAEPretraining(
    in_channels: int,
    *,
    embed_dim: int = 384,
    encoder_depth: int = 12,
    decoder_depth: int = 4,
    num_heads: int = 6,
    decoder_num_heads: int = 6,
    num_groups: int = 64,
    group_size: int = 32,
    encoder_local_channels: Sequence[int] = (128, 256),
    encoder_global_channels: Sequence[int] = (512,),
    pos_embed_channels: Sequence[int] = (128,),
    mask_ratio: float = 0.6,
    drop_path: float = 0.1,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = LayerNorm,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    spatial_dim: int = 3,
)

Bases: PretrainingModel

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

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

Parameters:

in_channels (int) –

Number of input channels (unused beyond coordinates; kept for the registry contract).
embed_dim (int, default: 384 ) –

Token channels.
encoder_depth (int, default: 12 ) –

Number of encoder transformer blocks.
decoder_depth (int, default: 4 ) –

Number of decoder transformer blocks.
num_heads (int, default: 6 ) –

Number of encoder attention heads.
decoder_num_heads (int, default: 6 ) –

Number of decoder attention heads.
num_groups (int, default: 64 ) –

Number of groups (FPS centers) per sample.
group_size (int, default: 32 ) –

Number of neighbors per group.
encoder_local_channels (Sequence[int], default: (128, 256) ) –

Hidden widths of the patch embedder's per-point MLP.
encoder_global_channels (Sequence[int], default: (512,) ) –

Hidden widths of the patch embedder's per-group MLP.
pos_embed_channels (Sequence[int], default: (128,) ) –

Hidden widths of the encoder and decoder positional-embedding MLPs.
mask_ratio (float, default: 0.6 ) –

Fraction of tokens to mask.
drop_path (float, default: 0.1 ) –

Stochastic-depth drop-path rate.
act (Union[str, Callable, None], default: 'gelu' ) –

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

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

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

Extra arguments for the normalization.
spatial_dim (int, default: 3 ) –

Spatial dimension of the input point cloud.

Shape

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

Methods:

configure_MAE_encoder –

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] ¶

configure_decoder_pos_embed() -> MLP

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

configure_MAE_decoder [source] ¶

configure_MAE_decoder() -> TransformerDecoder

Build the transformer decoder reconstructing the masked tokens.

configure_increase_dim [source] ¶

configure_increase_dim() -> MLP

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

point_mae

TransformerEncoder [source] ¶

TransformerDecoder [source] ¶

MaskTransformer [source] ¶

mask_center_rand [source] ¶

PointMAEClassification [source] ¶

num_features property ¶

configure_encoder [source] ¶

configure_pos_embed [source] ¶

configure_blocks [source] ¶

PointMAEPartSegmentation [source] ¶

num_features property ¶

configure_encoder [source] ¶

configure_pos_embed [source] ¶

configure_blocks [source] ¶

configure_label_conv [source] ¶

configure_propagation_0 [source] ¶

PointMAEPretraining [source] ¶

configure_MAE_encoder [source] ¶

configure_decoder_pos_embed [source] ¶

configure_MAE_decoder [source] ¶

configure_increase_dim [source] ¶

num_features `property` ¶

num_features `property` ¶