point_patch_embed

Mini-PointNet patch embedding turning local point groups into tokens.

Classes:

PointPatchEmbed –

Mini-PointNet patch (token) embedding for grouped point clouds.

PointPatchEmbed ¶

PointPatchEmbed(
    embed_dim: int,
    in_channels: int = 3,
    local_channels: Sequence[int] = (128, 256),
    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

Mini-PointNet patch (token) embedding for grouped point clouds.

Embeds each local group of points into a single token via a two-stage shared MLP with an intermediate per-group max-pool and global-feature concatenation, as used by the masked / autoregressive point self-supervised models (Point-MAE, Point-BERT, PointGPT, Point-M2AE). A shared \(1 \times 1\) convolution over \((B, C, M)\) is equivalent to a MLP over the feature dim, so both stages are plain PyG MLPs. local_mlp maps in_channels \(\to \text{local\_channels}\); the per-group max-pool is concatenated to give \(2 \cdot \text{local\_channels}[-1]\); global_mlp maps that to embed_dim through global_channels. The module is permutation-invariant over the points within a group.

Parameters:

embed_dim (int) –

Output token dimension.
in_channels (int, default: 3 ) –

Channels per input point (\(3\) for coordinates only, plus any concatenated features).
local_channels (Sequence[int], default: (128, 256) ) –

Hidden widths of the per-point MLP (in_channels is prepended).
global_channels (Sequence[int], default: (512,) ) –

Hidden widths of the per-group MLP (\(2 \cdot \text{local\_channels}[-1]\) input and embed_dim output are added).
act (Union[str, Callable, None], default: 'relu' ) –

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

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

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

Extra arguments for the normalization.

Shape

Input: \((B, G, M, C)\) where \(B\) is the batch size, \(G\) the number of groups, \(M\) the group size, and \(C\) = in_channels.
Output: \((B, G, D)\) where \(D\) = embed_dim.

Example

import torch
from torch_pointcloud.layers import PointPatchEmbed

embed = PointPatchEmbed(embed_dim=384)
tokens = embed(torch.randn(2, 64, 32, 3))
print(tokens.shape)