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

Classes:
-
EncoderBlockβA stack of transformer blocks that adds the positional embedding before every block.
-
ConvResBlock1dβA residual block used by the pre-training feature-propagation extraction.
-
FeaturePropagationβInterpolate features from a coarse set to a fine set, fuse and refine them.
-
HierarchicalEncoderβMulti-scale hierarchical transformer encoder of Point-M2AE.
-
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.
-
HierarchicalEncoderMAEβHierarchical encoder variant used for pre-training, with multi-scale back-propagated masking.
Functions:
-
local_attention_maskβConvert a boolean neighbor mask into the additive pre-softmax bias of the shared attention.
-
local_att_maskβCompute the boolean local-attention mask of a center set from a pairwise-distance threshold.
-
dense_centers_to_packedβFlatten a densified center set \((B, G, 3)\) into packed coordinates and a batch index.
-
multi_scale_groupβBuild the multi-scale grouping by repeatedly applying FPS + KNN on the previous-stage centers.
EncoderBlock [source] ΒΆ
EncoderBlock(
embed_dim: int,
depth: int,
num_heads: int,
mlp_ratio: float = 4.0,
qkv_bias: bool = False,
dropout: float = 0.0,
attn_drop_rate: float = 0.0,
drop_path: Union[float, List[float]] = 0.0,
act: Union[str, Callable, None] = "gelu",
norm: Union[str, Callable, None] = LayerNorm,
)
Bases: Module
A stack of transformer blocks that adds the positional embedding before every block.
Parameters:
-
embed_dim(int) βThe number of channels.
-
depth(int) βThe number of transformer blocks.
-
num_heads(int) βThe number of attention heads.
-
mlp_ratio(float, default:4.0) βThe feed-forward expansion ratio.
-
qkv_bias(bool, default:False) βWhether to add a bias to the query / key / value projection.
-
dropout(float, default:0.0) βThe dropout rate for the MLP and attention output projections.
-
attn_drop_rate(float, default:0.0) βThe attention dropout rate.
-
drop_path(Union[float, List[float]], default:0.0) βThe stochastic depth rate, either a scalar or a per-block list.
-
act(Union[str, Callable, None], default:'gelu') βThe activation function used in the feed-forward MLP.
-
norm(Union[str, Callable, None], default:LayerNorm) βThe normalization applied before attention and the MLP.
Shape
- Input: \((B, N, C)\)
- Output: \((B, N, C)\)
ConvResBlock1d [source] ΒΆ
ConvResBlock1d(
channels: int,
act: Union[str, Callable, None] = "gelu",
act_kwargs: Optional[Dict[str, Any]] = None,
norm: Union[str, Callable, None] = "batch_norm",
norm_kwargs: Optional[Dict[str, Any]] = None,
)
Bases: Module
A residual block used by the pre-training feature-propagation extraction.
Computes \(\text{act}(W_2(\text{act}(W_1 x)) + x)\) where each linear layer is followed by a normalization. The residual addition prevents expressing this as a single plain-last MLP, so the inner net1 (linear -> norm -> act) and net2 (linear -> norm) chains are built as MLP and the residual add is kept in forward. A shared \(1 \times 1\) convolution over \((B, C, N)\) is equivalent to a MLP over the flattened feature dim.
Parameters:
-
channels(int) βThe number of input and output channels.
-
act(Union[str, Callable, None], default:'gelu') βThe activation function.
-
act_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the activation.
-
norm(Union[str, Callable, None], default:'batch_norm') βThe normalization function.
-
norm_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the normalization.
Shape
- Input: \((B \cdot N, C)\)
- Output: \((B \cdot N, C)\)
FeaturePropagation [source] ΒΆ
FeaturePropagation(
in_channels: int,
out_channels: int,
blocks: int = 1,
act: Union[str, Callable, None] = "gelu",
act_kwargs: Optional[Dict[str, Any]] = None,
norm: Union[str, Callable, None] = "batch_norm",
norm_kwargs: Optional[Dict[str, Any]] = None,
)
Bases: Module
Interpolate features from a coarse set to a fine set, fuse and refine them.
Each fine point gathers the three nearest coarse points with the shared packed PointNet2FeaturePropagation interpolation, optionally concatenates the fine features, then applies a fuse MLP and a residual extraction stack. Used by the pre-training decoder. The interpolation matches the reference inverse-distance-squared weighting (\(1 / (d^2 + 10^{-8})\)) up to the \(\sim 10^{-6}\) difference between the direct \(\lVert a - b \rVert^2\) distance and the reference algebraic expansion.
Parameters:
-
in_channels(int) βThe number of input channels (fine + coarse features).
-
out_channels(int) βThe number of output channels.
-
blocks(int, default:1) βThe number of residual blocks in the extraction stack.
-
act(Union[str, Callable, None], default:'gelu') βThe activation function.
-
act_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the activation.
-
norm(Union[str, Callable, None], default:'batch_norm') βThe normalization function.
-
norm_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the normalization.
Shape
centers_fine: \((B, N, 3)\)centers_coarse: \((B, S, 3)\)points1: \((B, N, C_1)\) orNonepoints2: \((B, S, C_2)\)- Output: \((B, N, C_\text{out})\)
HierarchicalEncoder [source] ΒΆ
HierarchicalEncoder(
encoder_depths: Sequence[int],
encoder_dims: Sequence[int],
local_radius: Sequence[float],
num_heads: int,
drop_path: float = 0.1,
with_norms: bool = True,
in_channels: int = 0,
token_local_channels: Sequence[int] = (128, 256),
token_global_channels: Sequence[int] = (512,),
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: Module
Multi-scale hierarchical transformer encoder of Point-M2AE.
Each stage embeds (or merges) tokens, computes a local-attention mask from the stage centers and runs a transformer block stack. The forward pass operates on the visible (unmasked) tokens; in eval mode no masking is applied. The encoder consumes the precomputed multi-scale grouping (neighborhoods, centers and neighbor indices) so the same grouping can be shared with the decoder.
Parameters:
-
encoder_depths(Sequence[int]) βThe number of transformer blocks per stage.
-
encoder_dims(Sequence[int]) βThe channel width per stage.
-
local_radius(Sequence[float]) βThe local-attention radius per stage (non-positive disables the mask).
-
num_heads(int) βThe number of attention heads.
-
drop_path(float, default:0.1) βThe maximum stochastic depth rate (linearly scaled across blocks).
-
with_norms(bool, default:True) βWhether to apply a per-stage output LayerNorm (disabled for the ModelNet40 / ScanObjectNN finetune heads).
-
in_channels(int, default:0) βThe number of per-point feature channels concatenated to the coordinates at the first stage (\(0\) for coordinates only).
-
token_local_channels(Sequence[int], default:(128, 256)) βHidden widths of the first-stage token embedder's per-point MLP (later stages derive their widths from
encoder_dims). -
token_global_channels(Sequence[int], default:(512,)) βHidden widths of the first-stage token embedder's per-group MLP.
-
act(Union[str, Callable, None], default:'relu') βThe activation function used in the token embedders and transformer blocks.
-
act_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the activation.
-
norm(Union[str, Callable, None], default:'batch_norm') βThe normalization function used in the token embedders.
-
norm_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the normalization.
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.
HierarchicalEncoderMAE [source] ΒΆ
HierarchicalEncoderMAE(
encoder_depths: Sequence[int],
encoder_dims: Sequence[int],
local_radius: Sequence[float],
num_heads: int,
mask_ratio: float,
drop_path: float = 0.1,
token_local_channels: Sequence[int] = (128, 256),
token_global_channels: Sequence[int] = (512,),
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: Module
Hierarchical encoder variant used for pre-training, with multi-scale back-propagated masking.
The coarsest-stage mask is sampled and propagated to finer stages through the neighbor indices, so a fine token is visible only if it contributes to a visible coarse token. Visible tokens are packed per sample to the longest visible length and processed with a per-stage local-attention mask.
Parameters:
-
encoder_depths(Sequence[int]) βThe number of transformer blocks per stage.
-
encoder_dims(Sequence[int]) βThe channel width per stage.
-
local_radius(Sequence[float]) βThe local-attention radius per stage.
-
num_heads(int) βThe number of attention heads.
-
mask_ratio(float) βThe fraction of coarsest-stage tokens to mask.
-
drop_path(float, default:0.1) βThe maximum stochastic depth rate.
-
token_local_channels(Sequence[int], default:(128, 256)) βHidden widths of the first-stage token embedder's per-point MLP (later stages derive their widths from
encoder_dims). -
token_global_channels(Sequence[int], default:(512,)) βHidden widths of the first-stage token embedder's per-group MLP.
-
act(Union[str, Callable, None], default:'relu') βThe activation function used in the token embedders and transformer blocks.
-
act_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the activation.
-
norm(Union[str, Callable, None], default:'batch_norm') βThe normalization function used in the token embedders.
-
norm_kwargs(Optional[Dict[str, Any]], default:None) βExtra keyword arguments for the normalization.
Methods:
-
rand_maskβDraw a random mask hiding
mask_ratioof the coarsest-stage tokens, independently for every sample.
local_attention_mask [source] ΒΆ
Convert a boolean neighbor mask into the additive pre-softmax bias of the shared attention.
Entries that are True (a point pair outside the local radius, or a padded / cross-sample pair) are pushed to a large negative value so they vanish after the softmax. The result is broadcast over the attention heads.
Parameters:
-
mask(Tensor) βBoolean / float mask of shape \((B, N, N)\) where non-zero marks a forbidden pair.
Returns:
-
TensorβThe additive attention bias of shape \((B, 1, N, N)\).
Shape
mask: \((B, N, N)\)- Output: \((B, 1, N, N)\)
local_att_mask [source] ΒΆ
Compute the boolean local-attention mask of a center set from a pairwise-distance threshold.
A pair of centers is masked (True) when their Euclidean distance is at least radius. The pairwise-distance matrix is recomputed only when the cached dist does not match the current number of centers, so consecutive stages with the same token count share it.
Parameters:
-
pos(Tensor) βCenter coordinates of shape \((B, N, 3)\).
-
radius(float) βThe local-attention radius.
-
dist(OptTensor, default:None) βAn optional cached pairwise-distance matrix from a previous call.
Returns:
-
TensorβA tuple
(mask, dist)with the boolean mask of shape \((B, N, N)\) and the pairwise-distance -
Tensorβmatrix of shape \((B, N, N)\).
Shape
pos: \((B, N, 3)\)mask: \((B, N, N)\)dist: \((B, N, N)\)
dense_centers_to_packed [source] ΒΆ
Flatten a densified center set \((B, G, 3)\) into packed coordinates and a batch index.
Every sample carries the same number of centers \(G\), so the packed batch index is simply each sample id repeated \(G\) times.
Parameters:
-
centers(Tensor) βDensified center coordinates of shape \((B, G, 3)\).
Returns:
-
Tuple[Tensor, Tensor]βA tuple
(pos, batch)withposof shape \((B \cdot G, 3)\) andbatchof shape \((B \cdot G,)\).
Shape
centers: \((B, G, 3)\)pos: \((B \cdot G, 3)\)batch: \((B \cdot G,)\)
multi_scale_group [source] ΒΆ
multi_scale_group(
pos: Tensor,
batch: Tensor,
num_groups: Sequence[int],
group_sizes: Sequence[int],
random_start: bool = False,
) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
Build the multi-scale grouping by repeatedly applying FPS + KNN on the previous-stage centers.
Parameters:
-
pos(Tensor) βPacked coordinates of shape \((N, 3)\).
-
batch(Tensor) βPer-point batch index of shape \((N,)\).
-
num_groups(Sequence[int]) βThe number of centers per stage.
-
group_sizes(Sequence[int]) βThe neighborhood size per stage.
Returns:
-
List[Tensor]βA tuple
(neighborhoods, centers, idxs)of per-stage tensors.centers[i]is \((B, G_i, 3)\), -
List[Tensor]βneighborhoods[i]is \((B, G_i, k_i, 3)\) andidxs[i]is the flat neighbor index into stage \(i-1\).