Skip to content

Sonata

Sonata pretrained encoder and linear-probing segmentation model.

First page of Sonata: Self-Supervised Learning of Reliable Point Representations

2503.16429 · March 2025

Classes:

SonataSegmentation 

SonataSegmentation(
    in_channels: int,
    num_classes: int,
    serialization_orders: Sequence[SerializationOrder] = (
        "z",
        "z-trans",
        "hilbert",
        "hilbert-trans",
    ),
    shuffle_serialization_orders: bool = True,
    strides: Sequence[int] = (2, 2, 2, 2),
    encoder_depths: Sequence[int] = (3, 3, 3, 12, 3),
    encoder_channels: Sequence[int] = (
        48,
        96,
        192,
        384,
        512,
    ),
    encoder_num_heads: Sequence[int] = (3, 6, 12, 24, 32),
    encoder_patch_size: Sequence[int] = (
        1024,
        1024,
        1024,
        1024,
        1024,
    ),
    norm: Union[str, Callable] = "layer_norm",
    act: Union[str, Callable] = "gelu",
    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.3,
    attn_kind: AttentionKind = "default",
    use_flash_attn: bool = True,
    upcast_attn: bool = False,
    upcast_softmax: bool = False,
    dropout: float = 0.0,
    pooling: str = "grid",
    stem_type: str = "linear",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    legacy: bool = False,
)

Bases: SegmentationModel

Sonata linear-probing segmentation model.

This variant follows the segmentation demo from facebookresearch/sonata: the encoder features are unpooled through the saved pooling inverses, concatenated with each parent stage, then projected by a linear segmentation head.

Note

The default (and registered) configuration enables flash attention (use_flash_attn=True), which requires the flash-attn package 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 the forward on CPU.

Methods:

Attributes:

  • num_features (int) –

    Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

num_features property 

num_features: int

Channel count \(C\) entering the head: every encoder stage unpooled and concatenated.

configure_encoder 

configure_encoder() -> PointTransformerV3Encoder

Build the PointTransformerV3Encoder backbone.

forward 

forward(
    x: Tensor, pos_grid: Tensor, batch: Tensor
) -> Tensor

Forward pass.

Parameters:

  • x (Tensor) –

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

  • pos_grid (Tensor) –

    Integer voxel-grid coordinates of shape \((N, 3)\) (used by the encoder for Z-order / Hilbert serialization, not float positions).

  • batch (Tensor) –

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