SphereFormer

SphereFormer segmentation model.

First page of Spherical Transformer for LiDAR-based 3D Recognition

2303.12766 · March 2023

Classes:

WindowedRelPosAttention –

Block-diagonal windowed multi-head self-attention with contextual relative-position encoding.
SphereFormerBlock –

Pre-norm transformer block wrapping WindowedRelPosAttention with an MLP, as in the reference.
SphereFormerUBlock –

Recursive UNet block: sparse residual blocks + windowed attention, then a downsample/upsample branch.
SphereFormerSegmentation –

SphereFormer semantic-segmentation model, as described in the paper

Functions:

cart2sphere –

Map Cartesian coordinates to spherical coordinates \((\theta, \phi, r)\) (degrees, degrees, meters).
exponential_split –

Quantize the radial relative position \(r_q - r_k\) into a signed, exponentially-growing bin index.

WindowedRelPosAttention ¶

WindowedRelPosAttention(
    embed_dim: int,
    num_heads: int,
    window_size: Tensor,
    window_size_sphere: Tensor,
    quant_size: Tensor,
    quant_size_sphere: Tensor,
    radial_split_exponent: float = 0.0125,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
)

Bases: SparseModule

Block-diagonal windowed multi-head self-attention with contextual relative-position encoding.

Runs two attentions in parallel and concatenates their heads: the first half of the heads attend within cubic (Cartesian) windows, the second half within radial (spherical) windows. Within each window every point attends to all others; scores are \(q \cdot k\) plus a learnable relative-position bias on both query and key, softmax-normalized per query, and the value is augmented with its own relative-position bias before the weighted sum. The windowed attention is computed by the sptr CUDA kernel (sptr.sparse_self_attention with pe_type="contextual", rel_query=rel_key=rel_value=True), mirroring the reference SparseMultiheadSASphereConcat.

Parameters:

embed_dim (int) –

Token dimension.
num_heads (int) –

Number of attention heads (split evenly between cubic and spherical branches).
window_size (Tensor) –

Cubic window size, of shape \((3,)\).
window_size_sphere (Tensor) –

Spherical window size \((\theta, \phi, r)\), of shape \((3,)\).
quant_size (Tensor) –

Cubic relative-position quantization size, of shape \((3,)\).
quant_size_sphere (Tensor) –

Spherical relative-position quantization size, of shape \((3,)\).
radial_split_exponent (float, default: 0.0125 ) –

Base bin width for the radial exponential split.
qkv_bias (bool, default: True ) –

Whether the fused QKV projection uses a bias.
qk_scale (Optional[float], default: None ) –

Optional override for the \(1/\sqrt{d}\) attention scale.

SphereFormerBlock ¶

SphereFormerBlock(
    embed_dim: int,
    num_heads: int,
    window_size: Tensor,
    window_size_sphere: Tensor,
    quant_size: Tensor,
    quant_size_sphere: Tensor,
    radial_split_exponent: float = 0.0125,
    mlp_ratio: float = 4.0,
    drop_path: float = 0.0,
    qkv_bias: bool = True,
)

Bases: Module

Pre-norm transformer block wrapping WindowedRelPosAttention with an MLP, as in the reference.

Applies x = x + attn(LN(x)) then x = x + mlp(LN(x)), with a GELU MLP of ratio mlp_ratio and an optional stochastic-depth drop_path on each residual branch.

Parameters:

embed_dim (int) –

Token dimension.
num_heads (int) –

Number of attention heads.
window_size (Tensor) –

Cubic window size, of shape \((3,)\).
window_size_sphere (Tensor) –

Spherical window size, of shape \((3,)\).
quant_size (Tensor) –

Cubic relative-position quantization size, of shape \((3,)\).
quant_size_sphere (Tensor) –

Spherical relative-position quantization size, of shape \((3,)\).
radial_split_exponent (float, default: 0.0125 ) –

Base bin width for the radial exponential split.
mlp_ratio (float, default: 4.0 ) –

Hidden-dim multiplier for the MLP.
drop_path (float, default: 0.0 ) –

Stochastic-depth rate.
qkv_bias (bool, default: True ) –

Whether the QKV projection uses a bias.

SphereFormerUBlock ¶

SphereFormerUBlock(
    planes: Sequence[int],
    block_reps: int,
    window_size: Tensor,
    window_size_sphere: Tensor,
    quant_size: Tensor,
    quant_size_sphere: Tensor,
    head_dim: int = 16,
    window_size_scale: Tuple[float, float] = (2.0, 2.0),
    drop_path: Sequence[float] = (0.0,),
    radial_split_exponent: float = 0.0125,
    indice_key_id: int = 1,
    sphere_layers: Sequence[int] = (1, 2, 3, 4, 5),
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Recursive UNet block: sparse residual blocks + windowed attention, then a downsample/upsample branch.

Each level runs block_reps sparse residual blocks, an optional SphereFormerBlock, then (for non-leaf levels) a strided sparse convolution into the next-deeper SphereFormerUBlock, an inverse convolution back, a skip concatenation, and block_reps tail residual blocks. The cubic and spherical window sizes are scaled by window_size_scale at every deeper level, mirroring the reference.

Parameters:

planes (Sequence[int]) –

Channel count of this level and all deeper levels.
block_reps (int) –

Number of residual blocks before (and after) the recursive branch.
window_size (Tensor) –

Cubic window size at this level, of shape \((3,)\).
window_size_sphere (Tensor) –

Spherical window size at this level, of shape \((3,)\).
quant_size (Tensor) –

Cubic quantization size at this level, of shape \((3,)\).
quant_size_sphere (Tensor) –

Spherical quantization size at this level, of shape \((3,)\).
head_dim (int, default: 16 ) –

Per-head dimension (sets num_heads = planes[0] // head_dim).
window_size_scale (Tuple[float, float], default: (2.0, 2.0) ) –

Pair (cubic_scale, sphere_scale) applied per deeper level.
drop_path (Sequence[float], default: (0.0,) ) –

Per-level stochastic-depth rates (indexed by level).
radial_split_exponent (float, default: 0.0125 ) –

Base bin width for the radial exponential split.
indice_key_id (int, default: 1 ) –

spconv indice-key id for this level.
sphere_layers (Sequence[int], default: (1, 2, 3, 4, 5) ) –

Levels (indice_key_id) that get a SphereFormerBlock.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization layer name / callable.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra keyword arguments for the normalization layer.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.

SphereFormerSegmentation ¶

SphereFormerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    base_channels: int = 32,
    layers: Sequence[int] = (32, 64, 128, 256, 256),
    block_reps: int = 2,
    head_dim: int = 16,
    window_size: Sequence[float] = (0.3, 0.3, 0.3),
    window_size_sphere: Sequence[float] = (2.0, 2.0, 80.0),
    quant_size: Sequence[float] = (0.0125, 0.0125, 0.0125),
    quant_size_sphere: Sequence[float] = (
        2.0 / 24,
        2.0 / 24,
        80.0 / 24,
    ),
    window_size_scale: Tuple[float, float] = (2.0, 1.5),
    sphere_layers: Sequence[int] = (1, 2, 3, 4, 5),
    radial_split_exponent: float = 0.0125,
    drop_path: float = 0.0,
    min_spatial_shape: int = 128,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SegmentationModel

SphereFormer semantic-segmentation model, as described in the paper Spherical Transformer for LiDAR-based 3D Recognition.

A sparse-convolution UNet32 backbone with a cubic + radial windowed self-attention block at every stage. The windowed attention is computed by the sptr CUDA kernel (an optional dependency), as in the reference. Inputs follow the packed convention: point features x, integer voxel-grid coordinates pos_grid, real-valued coordinates pos (used by the attention), and a per-point batch index. The output is per-point class logits.

Parameters:

in_channels (int) –

Input feature channels.
num_classes (int) –

Number of semantic classes.
base_channels (int, default: 32 ) –

Stem / level-0 channel count \(m\).
layers (Sequence[int], default: (32, 64, 128, 256, 256) ) –

Per-level channel counts (length = number of UNet levels).
block_reps (int, default: 2 ) –

Residual blocks per level (before and after the recursive branch).
head_dim (int, default: 16 ) –

Per-head dimension for the windowed attention.
window_size (Sequence[float], default: (0.3, 0.3, 0.3) ) –

Base cubic window size (voxel_size * patch_size * window), of shape \((3,)\).
window_size_sphere (Sequence[float], default: (2.0, 2.0, 80.0) ) –

Base spherical window size \((\theta, \phi, r)\), of shape \((3,)\).
quant_size (Sequence[float], default: (0.0125, 0.0125, 0.0125) ) –

Base cubic quantization size, of shape \((3,)\).
quant_size_sphere (Sequence[float], default: (2.0 / 24, 2.0 / 24, 80.0 / 24) ) –

Base spherical quantization size, of shape \((3,)\).
window_size_scale (Tuple[float, float], default: (2.0, 1.5) ) –

Pair (cubic_scale, sphere_scale) applied per deeper level.
sphere_layers (Sequence[int], default: (1, 2, 3, 4, 5) ) –

Levels (1-indexed) that receive a windowed-attention block.
radial_split_exponent (float, default: 0.0125 ) –

Base bin width for the radial exponential split.
drop_path (float, default: 0.0 ) –

Maximum stochastic-depth rate (linearly spread across levels).
min_spatial_shape (int, default: 128 ) –

Per-axis lower bound on the inferred sparse spatial shape.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization layer name / callable.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra keyword arguments for the normalization layer.
act (Union[str, Callable, None], default: 'relu' ) –

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

Extra keyword arguments for the activation.

Example

>>> import torch
>>> from torch_pointcloud.models import create_model
>>> model = create_model("sphereformer.semantickitti", task="segmentation").eval()  # doctest: +SKIP
>>> pos = torch.rand(1000, 3) * 10  # doctest: +SKIP
>>> pos_grid = (pos / 0.05).floor().long()  # doctest: +SKIP
>>> x = torch.cat([pos, torch.rand(1000, 1)], dim=1)  # doctest: +SKIP
>>> batch = torch.zeros(1000, dtype=torch.long)  # doctest: +SKIP
>>> logits = model(x, pos, pos_grid, batch)  # doctest: +SKIP
>>> logits.shape  # doctest: +SKIP
torch.Size([1000, 19])

Methods:

configure_input_conv –

Build the submanifold stem convolution.
configure_unet –

Build the recursive sparse UNet with windowed attention at every level.
configure_output_layer –

Build the final normalization and activation applied before the head.

Attributes:

num_features (int) –

Channel count \(C\) of the full-resolution features entering the head.

num_features `property` ¶

num_features: int

Channel count \(C\) of the full-resolution features entering the head.

configure_input_conv ¶

configure_input_conv() -> Module

Build the submanifold stem convolution.

configure_unet ¶

configure_unet() -> SphereFormerUBlock

Build the recursive sparse UNet with windowed attention at every level.

configure_output_layer ¶

configure_output_layer() -> Module

Build the final normalization and activation applied before the head.

cart2sphere ¶

cart2sphere(pos: Tensor) -> Tensor

Map Cartesian coordinates to spherical coordinates \((\theta, \phi, r)\) (degrees, degrees, meters).

The azimuth \(\theta = \operatorname{atan2}(y, x)\) and polar angle \(\phi = \operatorname{atan2}(\sqrt{x^2+y^2}, z)\) are returned in degrees (with \(\theta\) shifted to \([0, 360)\)); \(r = \sqrt{x^2+y^2+z^2}\) is the radius.

Parameters:

pos (Tensor) –

Cartesian coordinates.

Returns:

Tensor –

The spherical coordinates \((\theta, \phi, r)\).

Shape

Input: \((N, 3)\)
Output: \((N, 3)\)

Example

>>> import torch
>>> from torch_pointcloud.models.sphereformer import cart2sphere
>>> sphere = cart2sphere(torch.randn(8, 3))
>>> sphere.shape
torch.Size([8, 3])

exponential_split ¶

exponential_split(
    pos: Tensor,
    index_query: Tensor,
    index_key: Tensor,
    relative_position_index: Tensor,
    radial_split_exponent: float = 0.0125,
    offset: int = 24,
) -> Tensor

Quantize the radial relative position \(r_q - r_k\) into a signed, exponentially-growing bin index.

Reproduces the reference radial split: bins are symmetric around \(0\), double in width every two steps ([0, a), [a, 2a), [2a, 4a), [4a, 6a), [6a, 10a), ... with \(a\) the base bin width), and the sign of \(r_q - r_k\) selects the positive or negative half. The returned index is shifted by offset (half the number of rows of the radial relative-position table, the reference quant_size_scale) so it indexes the table without going negative, and clamped to \([0, 2 \cdot \text{offset} - 1]\): radial gaps beyond the outermost bins fall into those bins instead of overflowing the table. The signed bin index overwrites the third (radial) column of relative_position_index in place, matching the split_func contract of the sptr kernel.

Parameters:

pos (Tensor) –

Spherical coordinates whose third column is the radius \(r\).
index_query (Tensor) –

Per-pair query indices.
index_key (Tensor) –

Per-pair key indices.
relative_position_index (Tensor) –

Per-pair, per-axis relative-position table indices whose radial column is replaced with the signed exponential bin index.
radial_split_exponent (float, default: 0.0125 ) –

Base bin width of the radial exponential split.
offset (int, default: 24 ) –

Non-negative shift applied to the signed bin index; the radial table has \(2 \cdot \text{offset}\) rows.

Returns:

Tensor –

The updated relative_position_index with its radial column set to the signed bin index.

Shape

pos: \((N, 3)\)
index_query, index_key: \((M,)\)
relative_position_index: \((M, 3)\)
Output: \((M, 3)\)

SphereFormer

WindowedRelPosAttention ¶

SphereFormerBlock ¶

SphereFormerUBlock ¶

SphereFormerSegmentation ¶

num_features property ¶

configure_input_conv ¶

configure_unet ¶

configure_output_layer ¶

cart2sphere ¶

exponential_split ¶

num_features `property` ¶