point_transformer_v3

Point Transformer V3 classification and segmentation models.

First page of Point Transformer V3: Simpler, Faster, Stronger

2312.10035 · December 2023

Classes:

SerializedFeaturesDict –

PooledFeaturesDict of a Point Transformer V3 stage, carrying the serialization the decoder blocks re-use.
PointTransformerV3Block –

Transformer block over serialized patches: an xCPE sparse-convolution residual, then pre-normed
PointTransformerV3EncoderBlock –

One encoder stage: an optional pooling downsampling, then depth PointTransformerV3Block units.
PointTransformerV3DecoderBlock –

One decoder stage: an optional upsampling onto the skip resolution, then depth PointTransformerV3Block units
PointTransformerV3Encoder –

Point Transformer V3 encoder backbone.
PointTransformerV3Decoder –

Point Transformer V3 decoder with skip connections.
PointTransformerV3Classification –

PyTorch implementation of the Point Transformer V3 model, as described in the paper
PointTransformerV3Segmentation –

PyTorch implementation of the Point Transformer V3 model for segmentation tasks.

Functions:

serialize –

Encode voxel-grid coordinates along one or more space-filling curves and sort the points by each code.
encoder_checkpoint_filter_fn –

Rename the keys of a released Sonata, Concerto or Utonia encoder checkpoint to the PointTransformerV3Encoder
segmentation_checkpoint_filter_fn –

Rename the keys of a released segmentation checkpoint (a backbone. encoder and decoder with a seg_head) to
pretrained_encoder_hparams –

The PointTransformerV3Encoder hparams shared by the released Sonata, Concerto and Utonia encoders: a linear

SerializedFeaturesDict [source] ¶

Bases: PooledFeaturesDict

PooledFeaturesDict of a Point Transformer V3 stage, carrying the serialization the decoder blocks re-use.

Attributes:

serialized_code (Tensor) –

Serialization codes of the stage's points, shape \((L, N_s)\), one row per order.
serialized_order (Tensor) –

Permutation sorting the points along each order, shape \((L, N_s)\).
serialized_inverse (Tensor) –

Inverse of serialized_order, shape \((L, N_s)\).

PointTransformerV3Block [source] ¶

PointTransformerV3Block(
    channels: int,
    num_heads: int,
    patch_size: int = 48,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.0,
    act: Union[str, Callable] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    cpe_indice_key: Optional[str] = None,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = True,
    upcast_softmax: bool = True,
    rope_base: float = 10.0,
    legacy: bool = False,
)

Bases: Module

Transformer block over serialized patches: an xCPE sparse-convolution residual, then pre-normed patch attention and an MLP.

PointTransformerV3EncoderBlock [source] ¶

PointTransformerV3EncoderBlock(
    channels: int,
    depth: int,
    num_heads: int,
    patch_size: int,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: ValueCollection[float] = 0.0,
    act: Union[str, Callable] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    cpe_indice_key: Optional[str] = None,
    downsample: Optional[Module] = None,
    serialization_orders: Optional[
        Sequence[SerializationOrder]
    ] = None,
    shuffle_serialization_orders: bool = False,
    rope_base: float = 10.0,
    legacy: bool = False,
)

Bases: Module

One encoder stage: an optional pooling downsampling, then depth PointTransformerV3Block units.

Consecutive blocks cycle through the available serialization orders, so each attends over a differently ordered patch partition. Grid pooling re-serializes the pooled cloud, while serialized pooling derives the coarser codes by bit-shifting the finer ones.

PointTransformerV3DecoderBlock [source] ¶

PointTransformerV3DecoderBlock(
    channels: int,
    depth: int,
    num_heads: int,
    patch_size: int,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: ValueCollection[float] = 0.0,
    act: Union[str, Callable] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    cpe_indice_key: Optional[str] = None,
    upsample: Optional[SerializedUpsample] = None,
    cpe_from_skip: bool = True,
    rope_base: float = 10.0,
    legacy: bool = False,
)

Bases: Module

One decoder stage: an optional upsampling onto the skip resolution, then depth PointTransformerV3Block units cycling through the skip's serialization orders.

cpe_from_skip picks the features the first block's xCPE convolves after upsampling: the projected skip branch, as serialized unpooling leaves it in the sparse tensor, or the summed output, as grid unpooling does.

PointTransformerV3Encoder [source] ¶

PointTransformerV3Encoder(
    in_channels: int = 6,
    serialization_orders: Sequence[SerializationOrder] = (
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (2, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        512,
    ),
    encoder_num_heads: Sequence[int] = (2, 4, 8, 16, 32),
    encoder_patch_size: Sequence[int] = (
        48,
        48,
        48,
        48,
        48,
    ),
    act: Union[str, Callable] = "gelu",
    norm: Union[str, Callable] = "batch_norm",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    mlp_ratio: float = 4,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    pooling: str = "serialized",
    stem_type: str = "sparse_conv",
    rope_base: float = 10.0,
    legacy: bool = False,
)

Bases: Module

Point Transformer V3 encoder backbone.

Encoder-only backbone for feature extraction from 3D point clouds. Supports both sparse convolution (PTV3 Mode 1) and linear (Sonata / Mode 2) embedding stems, and both serialized (code-space) and grid-based pooling.

Parameters:

in_channels (int, default: 6 ) –

Number of input channels.
serialization_orders (Sequence[SerializationOrder], default: ('hilbert', 'hilbert-trans') ) –

Serialization orders for attention.
shuffle_serialization_orders (bool, default: True ) –

Shuffle orders each forward pass.
strides (Sequence[int], default: (2, 2, 2, 2) ) –

Downsampling strides between encoder stages.
encoder_depths (Sequence[int], default: (2, 2, 2, 6, 2) ) –

Number of blocks per encoder stage.
encoder_channels (Sequence[int], default: (32, 64, 128, 256, 512) ) –

Feature channels per encoder stage.
encoder_num_heads (Sequence[int], default: (2, 4, 8, 16, 32) ) –

Attention heads per encoder stage.
encoder_patch_size (Sequence[int], default: (48, 48, 48, 48, 48) ) –

Patch size per encoder stage.
norm (Union[str, Callable], default: 'batch_norm' ) –

Normalization layer type.
act (Union[str, Callable], default: 'gelu' ) –

Activation function type.
mlp_ratio (float, default: 4 ) –

MLP expansion ratio.
qkv_bias (bool, default: True ) –

Use bias in QKV projection.
qk_scale (Optional[float], default: None ) –

Custom QK scaling factor.
attn_drop (float, default: 0.0 ) –

Attention dropout rate.
proj_drop (float, default: 0.0 ) –

Projection dropout rate.
drop_path (float, default: 0.3 ) –

Drop path rate.
attn_kind (AttentionKind, default: 'default' ) –

Attention variant: "default" (vanilla, PT-V3 / Sonata / Concerto), "rpe" (PT-V3 with learned relative position bias), or "rope" (Utonia, 3D rotary position embedding on Q, K). The "rope" variant requires the real-valued pos argument at forward time.
use_flash_attn (bool, default: True ) –

Use Flash Attention. The registered configurations construct with use_flash_attn=True, which requires flash-attn and a CUDA device; pass use_flash_attn=False to run without it (the xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run on CPU).
upcast_attn (bool, default: False ) –

Upcast attention to fp32.
upcast_softmax (bool, default: False ) –

Upcast softmax to fp32.
pooling (str, default: 'serialized' ) –

Pooling strategy: "serialized" (code-space bit-shift) or "grid" (grid-coordinate clustering).
stem_type (str, default: 'sparse_conv' ) –

How to embed raw features: "sparse_conv" (SubMConv3d stem) or "linear" (linear projection).
rope_base (float, default: 10.0 ) –

RoPE frequency base. Only used when attn_kind="rope".
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the activation factory.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the normalization factory.
bias (bool, default: True ) –

Whether the stem and blocks use learnable bias where applicable.
legacy (bool, default: False ) –

Reproduce the reference implementation's v1.5.1 block xCPE bug (the block output was not written back to the sparse tensor the next block convolves; fixed in v1.5.2). The released weights need legacy=True; leave False (default) for new training.

Inputs

x: Float tensor of shape \((N, \text{in\_channels})\). pos_grid: Int tensor of shape \((N, 3)\) with voxel-grid coordinates. batch: Long tensor of shape \((N,)\).

Outputs

Encoded features at the deepest encoder level.

Methods:

configure_stem –

Build the embedding stem, either a LinearBlock or a SubMConv3dBlock.
configure_blocks –

Build the PointTransformerV3EncoderBlock stages, giving every stage but the first a pooling downsampling.

Attributes:

embedding_dim (int) –

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

embedding_dim `property` ¶

embedding_dim: int

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

configure_stem [source] ¶

configure_stem(
    in_channels: int,
    out_channels: int,
    norm: Union[str, Callable],
    act: Union[str, Callable],
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    stem_type: str = "sparse_conv",
) -> Module

Build the embedding stem, either a LinearBlock or a SubMConv3dBlock.

configure_blocks [source] ¶

configure_blocks(
    depths: Sequence[int],
    channels: Sequence[int],
    num_heads: Sequence[int],
    patch_sizes: Sequence[int],
    strides: Sequence[int],
    mlp_ratio: float = 4.0,
    bias: bool = True,
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.0,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    pooling: str = "serialized",
    serialization_orders: Optional[
        Sequence[SerializationOrder]
    ] = None,
    shuffle_serialization_orders: bool = False,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    rope_base: float = 10.0,
    legacy: bool = False,
) -> ModuleList

Build the PointTransformerV3EncoderBlock stages, giving every stage but the first a pooling downsampling.

PointTransformerV3Decoder [source] ¶

PointTransformerV3Decoder(
    encoder_channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        512,
    ),
    decoder_depths: Sequence[int] = (2, 2, 2, 2),
    decoder_channels: Sequence[int] = (256, 128, 64, 64),
    decoder_num_heads: Sequence[int] = (16, 8, 4, 4),
    decoder_patch_size: Sequence[int] = (48, 48, 48, 48),
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    rope_base: float = 10.0,
    legacy: bool = False,
    pooling: str = "serialized",
)

Bases: Module

Point Transformer V3 decoder with skip connections.

Decoder backbone that upsamples encoder features using skip connections from intermediate encoder stages. Used for dense prediction tasks like semantic segmentation.

Parameters:

encoder_channels (Sequence[int], default: (32, 64, 128, 256, 512) ) –

Channel sequence from the encoder (needed to derive skip-connection channels).
decoder_depths (Sequence[int], default: (2, 2, 2, 2) ) –

Number of blocks per decoder stage.
decoder_channels (Sequence[int], default: (256, 128, 64, 64) ) –

Feature channels per decoder stage.
decoder_num_heads (Sequence[int], default: (16, 8, 4, 4) ) –

Attention heads per decoder stage.
decoder_patch_size (Sequence[int], default: (48, 48, 48, 48) ) –

Patch size per decoder stage.
norm (Union[str, Callable], default: 'batch_norm' ) –

Normalization layer type.
act (Union[str, Callable], default: 'gelu' ) –

Activation function type.
mlp_ratio (float, default: 4 ) –

MLP expansion ratio.
qkv_bias (bool, default: True ) –

Use bias in QKV projection.
qk_scale (Optional[float], default: None ) –

Custom QK scaling factor.
attn_drop (float, default: 0.0 ) –

Attention dropout rate.
proj_drop (float, default: 0.0 ) –

Projection dropout rate.
drop_path (float, default: 0.3 ) –

Drop path rate.
attn_kind (AttentionKind, default: 'default' ) –

Attention variant ("default", "rpe", or "rope").
use_flash_attn (bool, default: True ) –

Use Flash Attention.
upcast_attn (bool, default: False ) –

Upcast attention to fp32.
upcast_softmax (bool, default: False ) –

Upcast softmax to fp32.
rope_base (float, default: 10.0 ) –

RoPE frequency base. Only used when attn_kind="rope".
pooling (str, default: 'serialized' ) –

Pooling the encoder used between stages ("serialized" or "grid"). After grid pooling the block xCPE convolves the summed upsampling output, after serialized pooling the projected skip, as the reference unpooling modules leave them.

Inputs

x: Encoded features at the deepest encoder level. intermediates: List of dicts from the encoder, each containing skip features, positions, batch indices, serialization tensors, and pooling inverse indices.

Outputs

Decoded features at the shallowest decoder level.

Methods:

configure_blocks –

Build the PointTransformerV3DecoderBlock stages, giving every stage an upsampling onto its skip resolution.

Attributes:

out_channels (int) –

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

out_channels `property` ¶

out_channels: int

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

configure_blocks [source] ¶

configure_blocks(
    depths: Sequence[int],
    channels: Sequence[int],
    skip_channels: Sequence[int],
    num_heads: Sequence[int],
    patch_sizes: Sequence[int],
    mlp_ratio: float = 4.0,
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.0,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    rope_base: float = 10.0,
    legacy: bool = False,
    pooling: str = "serialized",
) -> ModuleList

Build the PointTransformerV3DecoderBlock stages, giving every stage an upsampling onto its skip resolution.

PointTransformerV3Classification [source] ¶

PointTransformerV3Classification(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (2, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        512,
    ),
    encoder_num_heads: Sequence[int] = (2, 4, 8, 16, 32),
    encoder_patch_size: Sequence[int] = (
        48,
        48,
        48,
        48,
        48,
    ),
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    rope_base: float = 10.0,
    dropout: float = 0.0,
    global_pool: PoolLike = "max",
    pooling: str = "serialized",
    stem_type: str = "sparse_conv",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
    pdnorm_conditions: Optional[Sequence[str]] = None,
    condition: Optional[str] = None,
)

Bases: ClassificationModel

PyTorch implementation of the Point Transformer V3 model, as described in the paper Point Transformer V3: Simpler, Faster, Stronger by Xiaoyang Wu, Li Jiang, Peng-Shuai Wang, Zhijian Liu, Xihui Liu, Yu Qiao, Wanli Ouyang, Tong He, Hengshuang Zhao.

This implementation is based on the original implementation from Pointcept.

Important

This model requires spconv to be installed. It is also recommended to install flash-attn for faster attention. The registered configurations construct with use_flash_attn=True, which requires flash-attn and a CUDA device; pass use_flash_attn=False to run without it. The xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run on CPU. In addition, it is recommended to install ocnn if you want to use more serialization orders.

Parameters:

in_channels (int) –

Number of input channels (corresponding to the number of features).
num_classes (int) –

Number of output classes.
serialization_orders (Sequence[SerializationOrder], default: ('hilbert', 'hilbert-trans') ) –

Serialization orders to use for the PointTransformerV3Encoder.
shuffle_serialization_orders (bool, default: True ) –

Whether to shuffle the serialization orders each step.
strides (Sequence[int], default: (2, 2, 2, 2) ) –

Downsampling strides between encoder stages.
encoder_depths (Sequence[int], default: (2, 2, 2, 6, 2) ) –

Number of encoder blocks per stage.
encoder_channels (Sequence[int], default: (32, 64, 128, 256, 512) ) –

Number of channels per stage.
encoder_num_heads (Sequence[int], default: (2, 4, 8, 16, 32) ) –

Number of attention heads per stage.
encoder_patch_size (Sequence[int], default: (48, 48, 48, 48, 48) ) –

Patch size per stage.
norm (Union[str, Callable], default: 'batch_norm' ) –

Normalization layer to use.
act (Union[str, Callable], default: 'gelu' ) –

Activation function to use.
mlp_ratio (float, default: 4 ) –

MLP hidden dimension ratio inside each block.
qkv_bias (bool, default: True ) –

Whether to use bias in the QKV projection.
qk_scale (Optional[float], default: None ) –

Scaling factor for the QK matrix.
attn_drop (float, default: 0.0 ) –

Dropout rate for the attention.
proj_drop (float, default: 0.0 ) –

Dropout rate for the output projection of each block.
drop_path (float, default: 0.3 ) –

Stochastic depth rate.
attn_kind (AttentionKind, default: 'default' ) –

Attention variant: "default", "rpe", or "rope". The "rope" variant requires the real-valued pos argument at forward time.
use_flash_attn (bool, default: True ) –

Whether to use flash attention.
upcast_attn (bool, default: False ) –

Whether to upcast the attention to fp32.
upcast_softmax (bool, default: False ) –

Whether to upcast the softmax in fp32.
rope_base (float, default: 10.0 ) –

RoPE frequency base. Only used when attn_kind="rope".
dropout (float, default: 0.0 ) –

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

How to pool point features to a batch-level vector ("max", "mean", etc.).
pooling (str, default: 'serialized' ) –

Pooling between encoder stages ("serialized" or "grid").
stem_type (str, default: 'sparse_conv' ) –

Encoder stem: "sparse_conv" or "linear".
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the activation factory.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the normalization factory.

Inputs

x: Float tensor of shape \((N, \text{in\_channels})\). pos_grid: Int tensor of shape \((N, 3)\) with voxel-grid coordinates. batch: Long tensor of shape \((N,)\). pos: Float tensor of shape \((N, 3)\) with metric coordinates. Required when attn_kind="rope".

Outputs

logits: Float tensor of shape \((N, \text{num\_classes})\).

Methods:

configure_encoder –

Build the PointTransformerV3Encoder backbone.
reset_classifier –

Resets the classification head with new parameters.
forward_head –

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

Forward pass of the Point Transformer V3 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() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

reset_classifier [source] ¶

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

Resets the classification head with new parameters.

Note

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

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_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{embedding\_dim})\).
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: OptTensor,
    pos_grid: Tensor,
    batch: Tensor,
    pos: OptTensor = None,
    *,
    condition: Union[str, Sequence[str], None] = None,
) -> Tensor

Forward pass of the Point Transformer V3 classification network.

Parameters:

x (OptTensor) –

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

Integer grid coordinates of shape \((N, 3)\). The encoder uses these to derive the Z-order / Hilbert serialization index, so they must be voxel indices, not float positions.
batch (Tensor) –

Batch indices for each point of shape \((N,)\).
pos (OptTensor, default: None ) –

Real-valued metric positions of shape \((N, 3)\). Required when attn_kind="rope"; ignored otherwise.
condition (Union[str, Sequence[str], None], default: None ) –

Optional per-batch condition selecting the PDNorm inner norms (keyword-only).

Returns:

Tensor –

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

PointTransformerV3Segmentation [source] ¶

PointTransformerV3Segmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "hilbert",
        "hilbert-trans",
    ),
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (2, 2, 2, 6, 2),
    encoder_channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        512,
    ),
    encoder_num_heads: Sequence[int] = (2, 4, 8, 16, 32),
    encoder_patch_size: Sequence[int] = (
        48,
        48,
        48,
        48,
        48,
    ),
    decoder_depths: Sequence[int] = (2, 2, 2, 2),
    decoder_channels: Sequence[int] = (256, 128, 64, 64),
    decoder_num_heads: Sequence[int] = (16, 8, 4, 4),
    decoder_patch_size: Sequence[int] = (48, 48, 48, 48),
    norm: Union[str, Callable] = "batch_norm",
    act: Union[str, Callable] = "gelu",
    mlp_ratio: float = 4,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.3,
    shuffle_serialization_orders: bool = True,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    rope_base: float = 10.0,
    dropout: float = 0.0,
    pooling: str = "serialized",
    stem_type: str = "sparse_conv",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
    pdnorm_conditions: Optional[Sequence[str]] = None,
    condition: Optional[str] = None,
)

Bases: SemanticSegmentationModel

PyTorch implementation of the Point Transformer V3 model for segmentation tasks.

Based on the paper Point Transformer V3: Simpler, Faster, Stronger by Xiaoyang Wu, Li Jiang, Peng-Shuai Wang, Zhijian Liu, Xihui Liu, Yu Qiao, Wanli Ouyang, Tong He, Hengshuang Zhao.

This segmentation variant uses an encoder-decoder architecture with skip connections.

Parameters:

in_channels (int) –

Number of input channels.
num_classes (int) –

Number of output classes for segmentation.
serialization_orders (Sequence[SerializationOrder], default: ('hilbert', 'hilbert-trans') ) –

Serialization orders to use for the encoder.
strides (Sequence[int], default: (2, 2, 2, 2) ) –

Strides for the downsampling operations.
encoder_depths (Sequence[int], default: (2, 2, 2, 6, 2) ) –

Number of encoder blocks for each stage.
encoder_channels (Sequence[int], default: (32, 64, 128, 256, 512) ) –

Number of channels for each encoder block.
encoder_num_heads (Sequence[int], default: (2, 4, 8, 16, 32) ) –

Number of attention heads for each encoder block.
encoder_patch_size (Sequence[int], default: (48, 48, 48, 48, 48) ) –

Patch size for each encoder block.
decoder_depths (Sequence[int], default: (2, 2, 2, 2) ) –

Number of decoder blocks for each stage.
decoder_channels (Sequence[int], default: (256, 128, 64, 64) ) –

Number of channels for each decoder block.
decoder_num_heads (Sequence[int], default: (16, 8, 4, 4) ) –

Number of attention heads for each decoder block.
decoder_patch_size (Sequence[int], default: (48, 48, 48, 48) ) –

Patch size for each decoder block.
norm (Union[str, Callable], default: 'batch_norm' ) –

Normalization layer to use.
act (Union[str, Callable], default: 'gelu' ) –

Activation function to use.
mlp_ratio (float, default: 4 ) –

Ratio of the hidden dimension to the input dimension.
qkv_bias (bool, default: True ) –

Whether to use bias in the QKV projection.
qk_scale (Optional[float], default: None ) –

Scaling factor for the QK matrix.
attn_drop (float, default: 0.0 ) –

Dropout rate for the attention.
proj_drop (float, default: 0.0 ) –

Dropout rate for the projection.
drop_path (float, default: 0.3 ) –

Dropout rate for the drop path.
shuffle_serialization_orders (bool, default: True ) –

Whether to shuffle the serialization orders.
attn_kind (AttentionKind, default: 'default' ) –

Attention variant: "default", "rpe", or "rope". The "rope" variant requires the real-valued pos argument at forward time.
rope_base (float, default: 10.0 ) –

RoPE frequency base. Only used when attn_kind="rope".
use_flash_attn (bool, default: True ) –

Whether to use flash attention. The registered configurations construct with use_flash_attn=True, which requires flash-attn and a CUDA device; pass use_flash_attn=False to run without it (the xCPE sparse convolution still needs a spconv build matching the device; the standard CUDA wheel cannot run on CPU).
upcast_attn (bool, default: False ) –

Whether to upcast the attention.
upcast_softmax (bool, default: False ) –

Whether to upcast the softmax.
dropout (float, default: 0.0 ) –

Dropout on the per-point logits.
pooling (str, default: 'serialized' ) –

Inter-stage pooling ("serialized" or "grid").
stem_type (str, default: 'sparse_conv' ) –

Encoder stem ("sparse_conv" or "linear").
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the activation factory.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the normalization factory.

Methods:

configure_encoder –

Build the PointTransformerV3Encoder backbone.
configure_decoder –

Build the PointTransformerV3Decoder upsampling the coarsest features back through the encoder skips.
reset_classifier –

Resets the segmentation head with new parameters.
forward –

Forward pass of the Point Transformer V3 segmentation network.

Attributes:

num_features (int) –

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

num_features `property` ¶

num_features: int

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

configure_encoder [source] ¶

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

configure_decoder [source] ¶

configure_decoder() -> PointTransformerV3Decoder

Build the PointTransformerV3Decoder upsampling the coarsest features back through the encoder skips.

reset_classifier [source] ¶

reset_classifier(num_classes: int, **kwargs: Any) -> None

Resets the segmentation head with new parameters.

Note

To set an empty segmentation head, use num_classes=0.

Parameters:

num_classes (int) –

Number of output classes.
**kwargs (Any, default: {} ) –

Additional keyword arguments to pass to the segmentation head.

forward [source] ¶

forward(
    x: Tensor,
    pos_grid: Tensor,
    batch: Tensor,
    pos: OptTensor = None,
    *,
    condition: Union[str, Sequence[str], None] = None,
) -> Tensor

Forward pass of the Point Transformer V3 segmentation network.

Parameters:

x (Tensor) –

Per-point features of shape \((N, C)\).
pos_grid (Tensor) –

Integer grid coordinates of shape \((N, 3)\) used for serialization.
batch (Tensor) –

Batch indices for each point of shape \((N,)\).
pos (OptTensor, default: None ) –

Real-valued metric positions of shape \((N, 3)\). Required when attn_kind="rope"; ignored otherwise.
condition (Union[str, Sequence[str], None], default: None ) –

Optional per-batch condition selecting the PDNorm inner norms (keyword-only).

Returns:

Tensor –

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

serialize [source] ¶

serialize(
    pos_grid: Tensor,
    batch: Tensor,
    orders: Sequence[SerializationOrder],
    shuffle: bool = False,
) -> Tuple[Tensor, Tensor, Tensor]

Encode voxel-grid coordinates along one or more space-filling curves and sort the points by each code.

Parameters:

pos_grid (Tensor) –

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

Per-point batch index of shape \((N,)\).
orders (Sequence[SerializationOrder]) –

The \(L\) space-filling curves to encode along, one code row per order.
shuffle (bool, default: False ) –

Whether to permute orders before encoding, so consecutive blocks pick different curves.

Returns:

Tensor –

The serialization codes, the permutation sorting the points by code, and its inverse, each of
Tensor –

shape \((L, N)\).

Raises:

ValueError –

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

encoder_checkpoint_filter_fn [source] ¶

encoder_checkpoint_filter_fn(
    state_dict: Dict[str, Any], model: Module
) -> Dict[str, Any]

Rename the keys of a released Sonata, Concerto or Utonia encoder checkpoint to the PointTransformerV3Encoder ones, as the pretrained_filter_fn of a WeightsDict.

The masked-modelling embedding.mask_token is dropped. When model holds the encoder as its encoder submodule (a linear-probing segmentation model), the encoder keys move under encoder. and the seg_head of the probe head file becomes head.

Parameters:

state_dict (Dict[str, Any]) –

State dict keyed as in the released checkpoint files.
model (Module) –

Model the state dict is loaded into.

Returns:

Dict[str, Any] –

The state dict keyed as model expects.

segmentation_checkpoint_filter_fn [source] ¶

segmentation_checkpoint_filter_fn(
    state_dict: Dict[str, Any], model: Module
) -> Dict[str, Any]

Rename the keys of a released segmentation checkpoint (a backbone. encoder and decoder with a seg_head) to the PointTransformerV3Segmentation ones, as the pretrained_filter_fn of a WeightsDict.

The released decoder names its stages after the encoder stage they upsample to, so the deepest runs first; model numbers them in execution order: dec{i} of n stages becomes decoder.blocks.{n - 1 - i}.

Parameters:

state_dict (Dict[str, Any]) –

State dict keyed as in the released checkpoint file.
model (Module) –

PointTransformerV3Segmentation the state dict is loaded into.

Returns:

Dict[str, Any] –

The state dict keyed as model expects.

pretrained_encoder_hparams [source] ¶

pretrained_encoder_hparams(
    encoder_depths: Sequence[int],
    encoder_channels: Sequence[int],
    encoder_num_heads: Sequence[int],
) -> Dict[str, Any]

The PointTransformerV3Encoder hparams shared by the released Sonata, Concerto and Utonia encoders: a linear stem on 9 input channels (position, color, normal), grid pooling and node-mode layer norm. The sizes differ by depth, width and head count only.

Parameters:

encoder_depths (Sequence[int]) –

Number of blocks per stage.
encoder_channels (Sequence[int]) –

Feature width per stage.
encoder_num_heads (Sequence[int]) –

Attention heads per stage.

Returns:

Dict[str, Any] –

Keyword arguments for PointTransformerV3Encoder.

point_transformer_v3

SerializedFeaturesDict [source] ¶

PointTransformerV3Block [source] ¶

PointTransformerV3EncoderBlock [source] ¶

PointTransformerV3DecoderBlock [source] ¶

PointTransformerV3Encoder [source] ¶

embedding_dim property ¶

configure_stem [source] ¶

configure_blocks [source] ¶

PointTransformerV3Decoder [source] ¶

out_channels property ¶

configure_blocks [source] ¶

PointTransformerV3Classification [source] ¶

num_features property ¶

configure_encoder [source] ¶

reset_classifier [source] ¶

forward_head [source] ¶

forward [source] ¶

PointTransformerV3Segmentation [source] ¶

num_features property ¶

configure_encoder [source] ¶

configure_decoder [source] ¶

reset_classifier [source] ¶

forward [source] ¶

serialize [source] ¶

encoder_checkpoint_filter_fn [source] ¶

segmentation_checkpoint_filter_fn [source] ¶

pretrained_encoder_hparams [source] ¶

embedding_dim `property` ¶

out_channels `property` ¶

num_features `property` ¶

num_features `property` ¶