Voxel-Mamba

Voxel-Mamba detection model.

First page of Voxel Mamba: Group-Free State Space Models for Point Cloud based 3D Object Detection

2406.10700 · June 2024

Classes:

DownSparse –

Downsampling stage (the reference's DownSp): optional strided sparse conv + residual subm blocks.
DSB –

Dual-scale State Space Models block of Voxel Mamba.
VoxelMambaBackbone –

Group-free Voxel Mamba sparse 3D backbone (the reference's Voxel_Mamba_Waymo).
CenterHeadOutput –

Raw dense center-head maps over the BEV feature grid.
SeparateHead –

Per-attribute conv head of the center head.
CenterHead –

Center-based detection head producing per-pixel heatmaps and box regressions.
VoxelMambaDetection –

Voxel Mamba: group-free state-space 3D object detector (packed point format).

Functions:

build_hilbert_template –

Build the flat Hilbert-curve lookup table used to serialize voxels (the reference's curve_template).
hilbert_serialize –

Per-scene voxel orderings along the Hilbert curve (the reference's get_hilbert_index_3d_mamba_lite).

DownSparse ¶

DownSparse(
    channels: int,
    kernel_size: int,
    stride: int,
    num_blocks: int,
    *,
    indice_key: str,
    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

Downsampling stage (the reference's DownSp): optional strided sparse conv + residual subm blocks.

Parameters:

channels (int) –

Channels (constant through the stage).
kernel_size (int) –

Kernel size of the leading strided conv.
stride (int) –

Stride of the leading conv; if \(1\) the leading conv is an identity.
num_blocks (int) –

Number of trailing residual subm blocks.
indice_key (str) –

Base indice key for the stage.
act (Union[str, Callable, None], default: 'relu' ) –

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

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

Normalization type or callable.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

DSB ¶

DSB(
    d_model: int,
    *,
    down_kernel_size: Sequence[int],
    down_stride: Sequence[int],
    num_down: Sequence[int],
    indice_key: str,
    downsample_rank: int,
    down_resolution: bool,
    norm_epsilon: float,
    rms_norm: bool,
    fused_add_norm: bool,
    residual_in_fp32: bool,
    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

Dual-scale State Space Models block of Voxel Mamba.

A voxel sparse tensor is encoded at a high-resolution and a downsampled scale. The high-resolution scale runs a backward (sequence-flipped) Mamba pass and the low-resolution scale a forward Mamba pass, each over the voxels serialized along a Hilbert curve (the group-free sequence). The low-resolution scale is fused back through an inverse (or submanifold) sparse conv plus both high-resolution skips.

Parameters:

d_model (int) –

Voxel feature channels.
down_kernel_size (Sequence[int]) –

Kernel size per scale (high, low).
down_stride (Sequence[int]) –

Stride per scale (high, low).
num_down (Sequence[int]) –

Residual-block count per scale (high, low).
indice_key (str) –

Base indice key for this block.
downsample_rank (int) –

Hilbert template rank for the low-resolution scale.
down_resolution (bool) –

If True, fuse with an inverse conv; else with a subm conv.
norm_epsilon (float) –

LayerNorm epsilon for the Mamba output norms.
rms_norm (bool) –

Use RMSNorm instead of nn.LayerNorm inside the Mamba blocks.
fused_add_norm (bool) –

Use the fused add+norm kernel inside the Mamba blocks.
residual_in_fp32 (bool) –

Keep the Mamba residual stream in fp32.
act (Union[str, Callable, None], default: 'relu' ) –

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

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

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

Extra normalization arguments.

VoxelMambaBackbone ¶

VoxelMambaBackbone(
    d_model: int,
    grid_size: Sequence[int],
    *,
    num_stage: Sequence[int] = (2, 2, 2),
    num_down: Sequence[Sequence[int]] = (
        (0, 1),
        (0, 1),
        (0, 1),
    ),
    down_stride: Sequence[Sequence[int]] = (
        (1, 1),
        (1, 2),
        (1, 4),
    ),
    down_kernel_size: Sequence[Sequence[int]] = (
        (3, 3),
        (3, 3),
        (3, 5),
    ),
    down_resolution: Sequence[bool] = (False, True, True),
    downsample_rank: Sequence[int] = (9, 8, 7),
    extra_down: int = 5,
    norm_epsilon: float = 1e-05,
    rms_norm: bool = True,
    fused_add_norm: bool = True,
    residual_in_fp32: bool = True,
)

Bases: Module

Group-free Voxel Mamba sparse 3D backbone (the reference's Voxel_Mamba_Waymo).

A stack of DSB blocks serializes voxels along multi-scale Hilbert curves and applies bidirectional Mamba (state-space) blocks, with periodic sparse downsampling of the height axis. There is no windowing or grouping: the whole scene is one sequence.

Parameters:

d_model (int) –

Voxel feature channels.
grid_size (Sequence[int]) –

Voxel grid extent \((n_x, n_y, n_z)\).
num_stage (Sequence[int], default: (2, 2, 2) ) –

Number of DSB blocks per stage.
num_down (Sequence[Sequence[int]], default: ((0, 1), (0, 1), (0, 1)) ) –

Per-stage residual-block counts for the two conv_encoder scales.
down_stride (Sequence[Sequence[int]], default: ((1, 1), (1, 2), (1, 4)) ) –

Per-stage strides for the two conv_encoder scales.
down_kernel_size (Sequence[Sequence[int]], default: ((3, 3), (3, 3), (3, 5)) ) –

Per-stage kernel sizes for the two conv_encoder scales.
down_resolution (Sequence[bool], default: (False, True, True) ) –

Per-stage flag selecting inverse-conv (vs subm) fusion.
downsample_rank (Sequence[int], default: (9, 8, 7) ) –

Per-stage Hilbert template rank for the low-resolution scale.
extra_down (int, default: 5 ) –

Block index after which the final height-compression conv runs.
norm_epsilon (float, default: 1e-05 ) –

LayerNorm epsilon for the Mamba output norms.
rms_norm (bool, default: True ) –

Use RMSNorm inside the Mamba blocks.
fused_add_norm (bool, default: True ) –

Use the fused add+norm kernel inside the Mamba blocks.
residual_in_fp32 (bool, default: True ) –

Keep the Mamba residual stream in fp32.

CenterHeadOutput ¶

Bases: TypedDict

Raw dense center-head maps over the BEV feature grid.

Attributes:

center (Tensor) –

Sub-cell BEV center offset, shape \((B, 2, H, W)\).
center_z (Tensor) –

Absolute box height, shape \((B, 1, H, W)\).
dim (Tensor) –

Log box size, shape \((B, 3, H, W)\).
rot (Tensor) –

\((\cos\theta, \sin\theta)\), shape \((B, 2, H, W)\).
iou (Tensor) –

IoU-rectification prediction in \([-1, 1]\), shape \((B, 1, H, W)\).
heatmap (Tensor) –

Per-class center logits, shape \((B, C, H, W)\).

SeparateHead ¶

SeparateHead(
    in_channels: int,
    num_classes: int,
    num_layers: int = 2,
    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,
    bias: bool = False,
)

Bases: Module

Per-attribute conv head of the center head.

One small stack of \(3\times3\) convs per box attribute, applied to the shared BEV features: center \((2)\), center_z \((1)\), dim \((3)\), rot \((2)\), iou \((1)\) and a class heatmap. The branch widths are fixed by the box parametrization; only the number of classes and the conv depth are configurable.

Parameters:

in_channels (int) –

Input channels.
num_classes (int) –

Number of classes predicted by the heatmap branch.
num_layers (int, default: 2 ) –

Number of convs per branch.
act (Union[str, Callable, None], default: 'relu' ) –

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

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

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

Extra normalization arguments.
bias (bool, default: False ) –

Whether the hidden convs carry a bias.

Shape

Input: \((B, C_\text{in}, H, W)\).
Output: dict of \((B, C_\text{attr}, H, W)\) tensors keyed by attribute.

CenterHead ¶

CenterHead(
    in_channels: int,
    num_classes: int,
    *,
    shared_conv_channels: int = 64,
    num_head_layers: int = 2,
    bn_eps: float = 0.001,
    bn_momentum: float = 0.01,
    bias: bool = False,
)

Bases: Module

Center-based detection head producing per-pixel heatmaps and box regressions.

A shared \(3\times3\) conv reduces the BEV features, then a SeparateHead regresses the per-attribute maps (one shared head over all classes here).

Reference implementation: open-mmlab/OpenPCDet (CenterHead).

Parameters:

in_channels (int) –

Channels of the input BEV feature map.
num_classes (int) –

Number of foreground classes.
shared_conv_channels (int, default: 64 ) –

Channels of the shared conv before the separate heads.
num_head_layers (int, default: 2 ) –

Number of convs per separate-head branch.
bn_eps (float, default: 0.001 ) –

BatchNorm epsilon.
bn_momentum (float, default: 0.01 ) –

BatchNorm momentum.
bias (bool, default: False ) –

Whether convs preceding a norm carry a bias.

VoxelMambaDetection ¶

VoxelMambaDetection(
    in_channels: int = 5,
    num_classes: int = 3,
    *,
    voxel_size: Sequence[float] = (0.32, 0.32, 0.1875),
    point_cloud_range: Sequence[float] = (
        -74.88,
        -74.88,
        -2.0,
        74.88,
        74.88,
        4.0,
    ),
    d_model: int = 128,
    vfe_num_filters: Sequence[int] = (128, 128),
    layer_nums: Sequence[int] = (1, 2, 2),
    layer_strides: Sequence[int] = (1, 2, 2),
    num_filters: Sequence[int] = (128, 128, 256),
    upsample_strides: Sequence[float] = (1, 2, 4),
    num_upsample_filters: Sequence[int] = (128, 128, 128),
    shared_conv_channels: int = 64,
    rms_norm: bool = True,
    fused_add_norm: bool = True,
    norm_epsilon: float = 1e-05,
)

Bases: DetectionModel

Voxel Mamba: group-free state-space 3D object detector (packed point format).

Reference: Zhang et al., 2024. Reference implementation: gwenzhang/Voxel-Mamba (built on DSVT).

Voxels are serialized into a single Hilbert-curve sequence and processed by bidirectional Mamba (state-space) blocks (no windowing / grouping), then scattered to a BEV map, refined by a 2D residual backbone, and decoded by a center-based head.

Parameters:

in_channels (int, default: 5 ) –

Raw point feature channels including xyz (e.g. \(5\) for Waymo).
num_classes (int, default: 3 ) –

Number of foreground classes.
voxel_size (Sequence[float], default: (0.32, 0.32, 0.1875) ) –

Voxel size \((v_x, v_y, v_z)\).
point_cloud_range (Sequence[float], default: (-74.88, -74.88, -2.0, 74.88, 74.88, 4.0) ) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
d_model (int, default: 128 ) –

Voxel feature channels of the Mamba backbone.
vfe_num_filters (Sequence[int], default: (128, 128) ) –

PFN widths of the dynamic mean VFE.
layer_nums (Sequence[int], default: (1, 2, 2) ) –

2D backbone residual-block counts per level.
layer_strides (Sequence[int], default: (1, 2, 2) ) –

2D backbone downsample strides per level.
num_filters (Sequence[int], default: (128, 128, 256) ) –

2D backbone channel widths per level.
upsample_strides (Sequence[float], default: (1, 2, 4) ) –

2D backbone upsample factors per level.
num_upsample_filters (Sequence[int], default: (128, 128, 128) ) –

2D backbone upsample channels per level.
shared_conv_channels (int, default: 64 ) –

Channels of the head's shared conv.
rms_norm (bool, default: True ) –

Use RMSNorm inside the Mamba blocks.
fused_add_norm (bool, default: True ) –

Use the fused add+norm kernel inside the Mamba blocks.
norm_epsilon (float, default: 1e-05 ) –

LayerNorm epsilon for the Mamba output norms.

Methods:

configure_vfe –

Build the dynamic mean voxel feature encoder.
configure_backbone_3d –

Build the Hilbert-serialized Mamba voxel backbone.
configure_backbone –

Build the residual 2D BEV backbone.
configure_head –

Build the center-based detection head.
decode –

Decode center-head predictions into raw candidate detections (no NMS).
reset_classifier –

Replace the classification branch of the detection head for num_classes outputs.

Attributes:

num_features (int) –

Channel count \(C\) of the BEV feature map entering the head.

num_features `property` ¶

num_features: int

Channel count \(C\) of the BEV feature map entering the head.

configure_vfe ¶

configure_vfe() -> DynamicMeanVFE

Build the dynamic mean voxel feature encoder.

configure_backbone_3d ¶

configure_backbone_3d() -> VoxelMambaBackbone

Build the Hilbert-serialized Mamba voxel backbone.

configure_backbone ¶

configure_backbone() -> BaseBEVResBackbone

Build the residual 2D BEV backbone.

configure_head ¶

configure_head() -> CenterHead

Build the center-based detection head.

decode ¶

decode(
    out: CenterHeadOutput,
    *,
    score_threshold: float = 0.0,
    top_k: int = 500,
    iou_rectifier: Sequence[float] = (0.68, 0.71, 0.65),
) -> Detection3D

Decode center-head predictions into raw candidate detections (no NMS).

Peaks of the (sigmoid) heatmap give candidate centers; box attributes are gathered at those peaks, mapped to world coordinates, and rescored by the predicted IoU (\(s^{1 - r_c} \cdot \text{iou}^{r_c}\) with a per-class rectifier \(r_c\), as in the reference). The full candidate set is returned; the evaluation pipeline applies score thresholding and per-class 3D NMS via the torch_pointcloud.utils.box3d utilities.

Parameters:

out (CenterHeadOutput) –

A CenterHeadOutput from forward.
score_threshold (float, default: 0.0 ) –

Minimum (pre-rectification) heatmap score to keep a peak; the non-filtering \(0\) default returns every peak (the reference protocol filters at \(0.1\)).
top_k (int, default: 500 ) –

Number of heatmap peaks gathered per scene.
iou_rectifier (Sequence[float], default: (0.68, 0.71, 0.65) ) –

Per-class IoU-rectification exponent, one entry per class (the default holds the reference Waymo 3-class values).

Returns:

Detection3D –

Packed candidate detections {"boxes": (K, 7), "scores": (K,), "labels": (K,), "batch": (K,)}
Detection3D –

(PyG layout).

reset_classifier ¶

reset_classifier(num_classes: int) -> None

Replace the classification branch of the detection head for num_classes outputs.

Models whose head is not rebuildable in isolation raise NotImplementedError.

build_hilbert_template ¶

build_hilbert_template(
    rank: int,
    z_max: int,
    device: Union[str, device] = "cpu",
) -> Tensor

Build the flat Hilbert-curve lookup table used to serialize voxels (the reference's curve_template).

A cube of side \(N = 2^\text{rank}\) is enumerated in \((z, y, x)\) order and each voxel is mapped to its Hilbert-curve position, then the table is truncated to \(N \cdot N \cdot z_\max\) entries (the curve only needs to cover voxels up to \(z_\max\) in the height axis). The table is indexed by the flat coordinate \(z \cdot N \cdot N + y \cdot N + x\) to read a voxel's position along the curve.

This reproduces the reference template (tools/hilbert_curves/create_hilbert_curve_template.py) bit-exactly via hilbert.encode, avoiding a 260 MB precomputed-weight download.

Parameters:

rank (int) –

Number of bits per dimension; the cube side is \(N = 2^\text{rank}\).
z_max (int) –

Height extent the curve must cover; the table keeps the first \(N \cdot N \cdot z_\max\) entries.
device (Union[str, device], default: 'cpu' ) –

Device the template is built on.

Returns:

Tensor –

A long tensor of shape \((N \cdot N \cdot z_\max,)\) giving each voxel's position on the curve.

Shape

Output: \((N \cdot N \cdot z_\max,)\)

Example

>>> template = build_hilbert_template(rank=7, z_max=9)
>>> template.shape
torch.Size([147456])

hilbert_serialize ¶

hilbert_serialize(
    template: Tensor,
    voxel_indices: Tensor,
    batch_size: int,
    rank: int,
    shift: int,
) -> Tuple[List[Tensor], List[Tensor]]

Per-scene voxel orderings along the Hilbert curve (the reference's get_hilbert_index_3d_mamba_lite).

Each voxel's flat coordinate (after a constant shift on every axis) indexes template to read its Hilbert position; sorting those positions within a scene yields the forward ordering, and sorting the forward ordering yields the inverse that scatters Mamba outputs back to voxel order.

Parameters:

template (Tensor) –

Flat Hilbert lookup table from build_hilbert_template, shape \((\cdot,)\).
voxel_indices (Tensor) –

Voxel coordinates \((N, 4)\) as \((\text{batch}, z, y, x)\).
batch_size (int) –

Number of scenes \(B\) in the batch.
rank (int) –

Rank the template was built with; the curve grid side is \(2^\text{rank}\).
shift (int) –

Constant offset added to \(z\), \(y\) and \(x\) before indexing the table.

Returns:

Tuple[List[Tensor], List[Tensor]] –

(forward, inverse), each a length-\(B\) list of long index tensors.

Shape

voxel_indices: \((N, 4)\)

Voxel-Mamba

DownSparse ¶

DSB ¶

VoxelMambaBackbone ¶

CenterHeadOutput ¶

SeparateHead ¶

CenterHead ¶

VoxelMambaDetection ¶

num_features property ¶

configure_vfe ¶

configure_backbone_3d ¶

configure_backbone ¶

configure_head ¶

decode ¶

reset_classifier ¶

build_hilbert_template ¶

hilbert_serialize ¶

num_features `property` ¶