Skip to content

lion

LION detection model.

First page of LION: Linear Group RNN for 3D Object Detection in Point Clouds

2407.18232 Β· July 2024

Classes:

  • BiMamba –

    Bidirectional Mamba mixer used by the LION linear group RNN (Mamba in LION's vendored ops).

  • MambaBlock –

    Post-norm residual wrapper around BiMamba (Block).

  • FlattenedWindowMapping –

    Window grouping / serialization for the linear group RNN (FlattenedWindowMapping).

  • PatchMerging3D –

    Voxel-generation / down-scaling step of LION (PatchMerging3D).

  • PatchExpanding3D –

    Scatter-back step pairing a PatchMerging3D (PatchExpanding3D).

  • LIONLayer –

    One linear group RNN layer: window-group serialize + Mamba per direction (LIONLayer).

  • LIONBlock –

    Hierarchical encoder/decoder stage of the LION backbone (LIONBlock).

  • LION3DBackbone –

    LION sparse 3D backbone: hierarchical linear group RNN over voxels (LION3DBackboneOneStride).

  • TransFusionHeadOutput –

    Raw per-query TransFusion head predictions plus the dense heatmap.

  • SeparateHeadTransfusion –

    Per-attribute prediction head of TransFusion (SeparateHead_Transfusion).

  • TransFusionDecoderLayer –

    Post-norm TransformerDecoderLayer over \((B, C, N)\) query and key maps (TransformerDecoderLayer).

  • TransFusionHead –

    Query-based TransFusion detection head (TransFusionHead).

  • LIONDetection –

    LION: linear group RNN (Mamba) 3D object detector (packed point format).

Functions:

  • window_partition –

    Assign each voxel to a 3D window and return its in-window offset (get_window_coors_shift_v2).

BiMamba [source] ΒΆ

BiMamba(
    d_model: int,
    d_state: int = 16,
    d_conv: int = 4,
    expand: int = 2,
)

Bases: Module

Bidirectional Mamba mixer used by the LION linear group RNN (Mamba in LION's vendored ops).

The input window-group sequence is projected to the Mamba inner width and scanned both forward and (sequence-flipped) backward by two independent selective-scan branches with separate input convs, \(\Delta\) projections and state matrices (\(A\), \(A_b\), \(D\), \(D_b\)); the two scans are summed (after un-flipping the backward branch) and projected back to \(d_\text{model}\). This is the operator the paper denotes as the linear RNN over each spatially grouped window.

Parameters:

  • d_model (int) –

    Feature channels of the mixer input/output.

  • d_state (int, default: 16 ) –

    SSM state width \(N\).

  • d_conv (int, default: 4 ) –

    Depthwise causal-conv kernel width.

  • expand (int, default: 2 ) –

    Inner-width expansion factor; \(d_\text{inner} = \text{expand} \cdot d_\text{model}\).

Shape
  • Input: \((B, L, d_\text{model})\)
  • Output: \((B, L, d_\text{model})\)

MambaBlock [source] ΒΆ

MambaBlock(
    d_model: int,
    d_state: int = 16,
    d_conv: int = 4,
    expand: int = 2,
)

Bases: Module

Post-norm residual wrapper around BiMamba (Block).

Computes \(x + \text{LayerNorm}(\text{BiMamba}(x))\) (post-norm residual, the LION default).

Parameters:

  • d_model (int) –

    Feature channels.

  • d_state (int, default: 16 ) –

    SSM state width passed to BiMamba.

  • d_conv (int, default: 4 ) –

    Causal-conv kernel width passed to BiMamba.

  • expand (int, default: 2 ) –

    Inner-width expansion passed to BiMamba.

FlattenedWindowMapping [source] ΒΆ

FlattenedWindowMapping(
    window_shape: Sequence[int],
    group_size: int,
    shift: bool,
)

Bases: Module

Window grouping / serialization for the linear group RNN (FlattenedWindowMapping).

Builds, for a voxel set, the index maps that (a) pad each scene up to a multiple of group_size (flat2win / win2flat) and (b) sort voxels into space-filling window order along the \(x\)- and \(y\)-major directions, so the Mamba operator can run over fixed-length contiguous groups.

Parameters:

  • window_shape (Sequence[int]) –

    Window size \((w_x, w_y, w_z)\).

  • group_size (int) –

    Sequence length of each Mamba group.

  • shift (bool) –

    Whether windows are shifted by half their size.

PatchMerging3D [source] ΒΆ

PatchMerging3D(
    dim: int,
    out_dim: int = -1,
    down_scale: Sequence[int] = (2, 2, 2),
    diffusion: bool = False,
    diff_scale: float = 0.2,
)

Bases: Module

Voxel-generation / down-scaling step of LION (PatchMerging3D).

A submanifold conv (LayerNorm + GELU) refines features, then voxels are merged onto a coarser grid by summing features that fall into the same down-scaled cell. When diffusion is set, the top diff_scale fraction of voxels (ranked by mean activation) are densified by spawning zero-feature neighbors before merging, the 3D voxel-generation that lets the linear RNN reach empty space.

Parameters:

  • dim (int) –

    Input feature channels.

  • out_dim (int, default: -1 ) –

    Output channels of the post-merge norm; -1 keeps dim.

  • down_scale (Sequence[int], default: (2, 2, 2) ) –

    Per-axis merge factor \((s_x, s_y, s_z)\).

  • diffusion (bool, default: False ) –

    Enable the voxel-generation (densification) step.

  • diff_scale (float, default: 0.2 ) –

    Fraction of voxels expanded when diffusion is set.

PatchExpanding3D [source] ΒΆ

PatchExpanding3D(dim: int)

Bases: Module

Scatter-back step pairing a PatchMerging3D (PatchExpanding3D).

Gathers the merged features back to the finer voxel layout (via the merge inverse map) and adds them onto the skip features at that scale.

Parameters:

  • dim (int) –

    Feature channels (unused; kept to mirror the reference signature).

LIONLayer [source] ΒΆ

LIONLayer(
    dim: int,
    window_shape: Sequence[int],
    group_size: int,
    direction: Sequence[str],
    shift: bool,
    d_state: int,
    d_conv: int,
    expand: int,
)

Bases: Module

One linear group RNN layer: window-group serialize + Mamba per direction (LIONLayer).

For each serialization direction (x, y) the voxels are gathered into fixed-length groups by FlattenedWindowMapping, passed through a MambaBlock, and scattered back to voxel order.

Parameters:

  • dim (int) –

    Feature channels.

  • window_shape (Sequence[int]) –

    Window size \((w_x, w_y, w_z)\).

  • group_size (int) –

    Mamba group length.

  • direction (Sequence[str]) –

    Serialization directions (e.g. ("x", "y")).

  • shift (bool) –

    Whether windows are shifted by half their size.

  • d_state (int) –

    SSM state width.

  • d_conv (int) –

    Causal-conv kernel width.

  • expand (int) –

    Inner-width expansion factor.

LIONBlock [source] ΒΆ

LIONBlock(
    dim: int,
    depth: int,
    down_scales: Sequence[Sequence[int]],
    window_shape: Sequence[int],
    group_size: int,
    direction: Sequence[str],
    shift: bool,
    d_state: int,
    d_conv: int,
    expand: int,
)

Bases: Module

Hierarchical encoder/decoder stage of the LION backbone (LIONBlock).

A depth-deep stack alternating LIONLayer (with a learned position embedding) and a down-scaling PatchMerging3D on the way down, then a matching decoder that scatters features back with PatchExpanding3D.

Parameters:

  • dim (int) –

    Feature channels.

  • depth (int) –

    Encoder/decoder depth.

  • down_scales (Sequence[Sequence[int]]) –

    Per-level merge factor \((s_x, s_y, s_z)\).

  • window_shape (Sequence[int]) –

    Window size \((w_x, w_y, w_z)\).

  • group_size (int) –

    Mamba group length.

  • direction (Sequence[str]) –

    Serialization directions.

  • shift (bool) –

    Whether the second layer per level uses shifted windows.

  • d_state (int) –

    SSM state width.

  • d_conv (int) –

    Causal-conv kernel width.

  • expand (int) –

    Inner-width expansion factor.

LION3DBackbone [source] ΒΆ

LION3DBackbone(
    grid_size: Sequence[int],
    *,
    channels: int = 128,
    num_layers: int = 4,
    depths: Sequence[int] = (2, 2, 2, 2),
    layer_down_scales: Sequence[Sequence[Sequence[int]]] = (
        ((2, 2, 2), (2, 2, 2)),
        ((2, 2, 2), (2, 2, 2)),
        ((2, 2, 2), (2, 2, 2)),
        ((2, 2, 2), (2, 2, 2)),
    ),
    window_shape: Sequence[Sequence[int]] = (
        (13, 13, 32),
        (13, 13, 16),
        (13, 13, 8),
        (13, 13, 4),
    ),
    group_size: Sequence[int] = (4096, 2048, 1024, 512),
    direction: Sequence[str] = ("x", "y"),
    diffusion: bool = True,
    diff_scale: float = 0.2,
    shift: bool = True,
    d_state: int = 16,
    d_conv: int = 4,
    expand: int = 2,
)

Bases: Module

LION sparse 3D backbone: hierarchical linear group RNN over voxels (LION3DBackboneOneStride).

Four LIONBlock stages, each followed by a height down-scaling PatchMerging3D (diffusion enabled), end in a single LIONLayer over the compressed grid. The output keeps the planar resolution (one-stride) and a height of 2 for BEV folding.

Parameters:

  • grid_size (Sequence[int]) –

    Voxel grid extent \((n_x, n_y, n_z)\).

  • channels (int, default: 128 ) –

    Backbone feature channels.

  • num_layers (int, default: 4 ) –

    Number of LIONBlock stages.

  • depths (Sequence[int], default: (2, 2, 2, 2) ) –

    Per-stage encoder/decoder depth.

  • layer_down_scales (Sequence[Sequence[Sequence[int]]], default: (((2, 2, 2), (2, 2, 2)), ((2, 2, 2), (2, 2, 2)), ((2, 2, 2), (2, 2, 2)), ((2, 2, 2), (2, 2, 2))) ) –

    Per-stage, per-depth merge factors.

  • window_shape (Sequence[Sequence[int]], default: ((13, 13, 32), (13, 13, 16), (13, 13, 8), (13, 13, 4)) ) –

    Per-stage window size \((w_x, w_y, w_z)\).

  • group_size (Sequence[int], default: (4096, 2048, 1024, 512) ) –

    Per-stage Mamba group length.

  • direction (Sequence[str], default: ('x', 'y') ) –

    Serialization directions.

  • diffusion (bool, default: True ) –

    Enable voxel-generation in the height-merge steps.

  • diff_scale (float, default: 0.2 ) –

    Fraction of voxels expanded by diffusion.

  • shift (bool, default: True ) –

    Whether shifted windows are used.

  • d_state (int, default: 16 ) –

    SSM state width.

  • d_conv (int, default: 4 ) –

    Causal-conv kernel width.

  • expand (int, default: 2 ) –

    Inner-width expansion factor.

TransFusionHeadOutput [source] ΒΆ

Bases: TypedDict

Raw per-query TransFusion head predictions plus the dense heatmap.

Attributes:

  • center (Tensor) –

    BEV center in feature-map cells, shape \((B, 2, Q)\).

  • height (Tensor) –

    Absolute box height, shape \((B, 1, Q)\).

  • dim (Tensor) –

    Log box size, shape \((B, 3, Q)\).

  • rot (Tensor) –

    \((\sin\theta, \cos\theta)\), shape \((B, 2, Q)\).

  • vel (Tensor) –

    BEV velocity, shape \((B, 2, Q)\).

  • iou (Tensor) –

    IoU-rectification prediction in \([-1, 1]\), shape \((B, 1, Q)\).

  • heatmap (Tensor) –

    Per-query class logits, shape \((B, C, Q)\).

  • query_heatmap_score (Tensor) –

    Dense-heatmap score gathered at each query cell, shape \((B, C, Q)\).

  • query_labels (Tensor) –

    Initial class of each query, shape \((B, Q)\).

  • dense_heatmap (Tensor) –

    Dense BEV class logits, shape \((B, C, W, H)\).

SeparateHeadTransfusion [source] ΒΆ

SeparateHeadTransfusion(
    in_channels: int,
    head_channels: int,
    num_classes: int,
    num_layers: int = 2,
    num_heatmap_layers: int = 2,
    init_bias: float = -2.19,
    bias: bool = False,
)

Bases: Module

Per-attribute prediction head of TransFusion (SeparateHead_Transfusion).

One small MLP per box attribute, applied to the per-query features: center \((2)\), height \((1)\), dim \((3)\), rot \((2)\), vel \((2)\), iou \((1)\) and a class heatmap. The reference's \(1 \times 1\) convolutions over \((B, C, Q)\) are equivalent to linear layers over the flattened query dim, so each branch is a plain MLP. The branch widths are fixed by the box parametrization; only the number of classes and the depths are configurable.

Parameters:

  • in_channels (int) –

    Input feature channels.

  • head_channels (int) –

    Hidden channels of the per-attribute MLP.

  • num_classes (int) –

    Number of classes predicted by the heatmap branch.

  • num_layers (int, default: 2 ) –

    Number of layers per box-attribute branch.

  • num_heatmap_layers (int, default: 2 ) –

    Number of layers of the heatmap branch.

  • init_bias (float, default: -2.19 ) –

    Bias initialization for the heatmap output layer.

  • bias (bool, default: False ) –

    Whether hidden layers carry a bias.

Shape
  • Input: \((B, C_\text{in}, Q)\) per-query features.
  • Output: dict of \((B, C_\text{attr}, Q)\) tensors keyed by attribute.

TransFusionDecoderLayer [source] ΒΆ

TransFusionDecoderLayer(
    embed_dim: int,
    num_heads: int,
    mlp_dim: int,
    dropout: float,
    activation: str,
    self_posembed: Module,
    cross_posembed: Module,
)

Bases: Module

Post-norm TransformerDecoderLayer over \((B, C, N)\) query and key maps (TransformerDecoderLayer).

The flattened query and key positions are embedded by the layer's own MLPs and added to the attention queries, keys and values.

Parameters:

  • embed_dim (int) –

    Model channels.

  • num_heads (int) –

    Number of attention heads.

  • mlp_dim (int) –

    FFN hidden width.

  • dropout (float) –

    Dropout probability.

  • activation (str) –

    FFN activation (relu/gelu).

  • self_posembed (Module) –

    Position embedding applied to the flattened query positions (self-attention).

  • cross_posembed (Module) –

    Position embedding applied to the flattened key positions (cross-attention).

TransFusionHead [source] ΒΆ

TransFusionHead(
    input_channels: int,
    num_classes: int,
    grid_size: Sequence[int],
    point_cloud_range: Sequence[float],
    voxel_size: Sequence[float],
    *,
    feature_map_stride: int = 2,
    hidden_channel: int = 128,
    num_proposals: int = 200,
    num_heads: int = 8,
    nms_kernel_size: int = 3,
    ffn_channel: int = 256,
    dropout: float = 0.0,
    bn_momentum: float = 0.1,
    activation: str = "relu",
    num_heatmap_layers: int = 2,
    query_radius: int = 20,
    iou_rectifier: float = 0.5,
    nms_radius: float = 0.175,
    local_max_classes: Sequence[int] = (),
    post_center_range: Sequence[float] = (
        -61.2,
        -61.2,
        -10.0,
        61.2,
        61.2,
        10.0,
    ),
)

Bases: Module

Query-based TransFusion detection head (TransFusionHead).

A shared conv produces a BEV feature map and a dense class heatmap; the top num_proposals heatmap peaks initialize object queries that attend (self- then local-cross-attention) to BEV features via a TransformerDecoderLayer. Per-query box attributes are regressed by a SeparateHeadTransfusion. decode rescores by IoU and applies per-task circular NMS (nuScenes ped/cone) to yield final boxes.

Reference implementation: mit-han-lab/bevfusion (TransFusionHead).

Parameters:

  • input_channels (int) –

    BEV feature channels feeding the shared conv.

  • num_classes (int) –

    Number of foreground classes.

  • grid_size (Sequence[int]) –

    Voxel grid extent \((n_x, n_y, n_z)\).

  • point_cloud_range (Sequence[float]) –

    Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).

  • voxel_size (Sequence[float]) –

    Voxel size \((v_x, v_y, v_z)\).

  • feature_map_stride (int, default: 2 ) –

    BEV stride relating feature pixels to metric coordinates.

  • hidden_channel (int, default: 128 ) –

    Query / decoder feature width.

  • num_proposals (int, default: 200 ) –

    Number of object queries.

  • num_heads (int, default: 8 ) –

    Decoder attention heads.

  • nms_kernel_size (int, default: 3 ) –

    Heatmap local-max pooling kernel.

  • ffn_channel (int, default: 256 ) –

    Decoder FFN width.

  • dropout (float, default: 0.0 ) –

    Decoder dropout.

  • bn_momentum (float, default: 0.1 ) –

    BatchNorm momentum override.

  • activation (str, default: 'relu' ) –

    Decoder FFN activation.

  • num_heatmap_layers (int, default: 2 ) –

    Layers in the heatmap branch of the prediction head.

  • query_radius (int, default: 20 ) –

    Half-width of the local cross-attention window.

  • iou_rectifier (float, default: 0.5 ) –

    Per-class exponent blending heatmap score with predicted IoU.

  • nms_radius (float, default: 0.175 ) –

    Per-task circular-NMS radius (only the local_max_classes use \(> 0\)).

  • local_max_classes (Sequence[int], default: () ) –

    Crowded small-object class indices (nuScenes pedestrian / traffic-cone): their heatmap peaks skip the kernel local-max NMS in predict and each gets its own circular NMS task in decode. Empty by default so the head stays agnostic to the label set.

  • post_center_range (Sequence[float], default: (-61.2, -61.2, -10.0, 61.2, 61.2, 10.0) ) –

    Box-center range filter applied at decode.

Methods:

  • reset_parameters –

    Start the weight matrices of the transformer decoder Xavier-uniform.

  • predict –

    Runs the head on a BEV feature map: initializes the queries from the dense heatmap, then decodes them.

  • decode –

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

reset_parameters [source] ΒΆ

reset_parameters() -> None

Start the weight matrices of the transformer decoder Xavier-uniform.

predict [source] ΒΆ

predict(inputs: Tensor) -> TransFusionHeadOutput

Runs the head on a BEV feature map: initializes the queries from the dense heatmap, then decodes them.

Parameters:

  • inputs (Tensor) –

    BEV features, shape \((B, C, W, H)\).

Returns:

decode [source] ΒΆ

decode(preds_dicts: TransFusionHeadOutput) -> Detection3D

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

Multiplies the sigmoid query scores by the gathered dense-heatmap score, recovers oriented boxes, rescores each by predicted IoU (iou_rectifier), and filters by the post-center range. The predicted BEV velocity \((v_x, v_y)\) of each kept box is returned under velocity. The full candidate set is returned; the evaluation pipeline applies the per-task circular NMS (on the local_max_classes, e.g. the nuScenes pedestrian / traffic-cone) via the torch_pointcloud.ops.box3d utilities (see the benchmark example).

Parameters:

Returns:

  • Detection3D –

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

  • Detection3D –

    (PyG layout), plus "velocity" \((K, 2)\).

LIONDetection [source] ΒΆ

LIONDetection(
    in_channels: int = 5,
    num_classes: int = 10,
    *,
    voxel_size: Sequence[float] = (0.3, 0.3, 0.25),
    point_cloud_range: Sequence[float] = (
        -54.0,
        -54.0,
        -5.0,
        54.0,
        54.0,
        3.0,
    ),
    channels: int = 128,
    vfe_num_filters: Sequence[int] = (128, 128),
    depths: Sequence[int] = (2, 2, 2, 2),
    window_shape: Sequence[Sequence[int]] = (
        (13, 13, 32),
        (13, 13, 16),
        (13, 13, 8),
        (13, 13, 4),
    ),
    group_size: Sequence[int] = (4096, 2048, 1024, 512),
    diffusion: bool = True,
    diff_scale: float = 0.2,
    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] = (0.5, 1, 2),
    num_upsample_filters: Sequence[int] = (128, 128, 128),
    feature_map_stride: int = 2,
    local_max_classes: Sequence[int] = (),
    d_state: int = 16,
    d_conv: int = 4,
    expand: int = 2,
)

Bases: DetectionModel

LION: linear group RNN (Mamba) 3D object detector (packed point format).

Reference: LION: Linear Group RNN for 3D Object Detection in Point Clouds (Liu et al., 2024). Reference implementation: happinesslz/LION.

Points are encoded into voxels by a dynamic mean VFE, processed by a hierarchical sparse backbone that serializes voxels into spatially grouped windows and runs a bidirectional Mamba operator (the linear group RNN), with periodic 3D voxel-generation / height-merging steps. The one-stride output is folded to a dense BEV map, refined by a residual 2D backbone, and decoded by a query-based TransFusionHead.

Parameters:

  • in_channels (int, default: 5 ) –

    Raw point feature channels including xyz (5 for nuScenes \(x, y, z, \text{intensity}, \Delta t\)).

  • num_classes (int, default: 10 ) –

    Number of foreground classes (10 for nuScenes).

  • voxel_size (Sequence[float], default: (0.3, 0.3, 0.25) ) –

    Voxel size \((v_x, v_y, v_z)\).

  • point_cloud_range (Sequence[float], default: (-54.0, -54.0, -5.0, 54.0, 54.0, 3.0) ) –

    Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).

  • channels (int, default: 128 ) –

    Backbone / VFE feature channels.

  • vfe_num_filters (Sequence[int], default: (128, 128) ) –

    PFN widths of the dynamic mean VFE.

  • depths (Sequence[int], default: (2, 2, 2, 2) ) –

    Per-stage encoder/decoder depth of the 3D backbone.

  • window_shape (Sequence[Sequence[int]], default: ((13, 13, 32), (13, 13, 16), (13, 13, 8), (13, 13, 4)) ) –

    Per-stage window size \((w_x, w_y, w_z)\).

  • group_size (Sequence[int], default: (4096, 2048, 1024, 512) ) –

    Per-stage Mamba group length.

  • diffusion (bool, default: True ) –

    Enable voxel-generation in the height-merge steps.

  • diff_scale (float, default: 0.2 ) –

    Fraction of voxels expanded by diffusion.

  • 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: (0.5, 1, 2) ) –

    2D backbone upsample factors per level.

  • num_upsample_filters (Sequence[int], default: (128, 128, 128) ) –

    2D backbone upsample channels per level.

  • feature_map_stride (int, default: 2 ) –

    BEV stride of the head.

  • local_max_classes (Sequence[int], default: () ) –

    Crowded small-object class indices passed to the head (nuScenes pedestrian / traffic-cone); see TransFusionHead.

  • d_state (int, default: 16 ) –

    SSM state width.

  • d_conv (int, default: 4 ) –

    Causal-conv kernel width.

  • expand (int, default: 2 ) –

    Inner-width expansion factor of the Mamba operator.

Methods:

  • configure_vfe –

    Build the dynamic mean voxel feature encoder.

  • configure_backbone_3d –

    Build the sparse 3D linear group RNN backbone.

  • configure_backbone –

    Build the residual 2D BEV backbone.

  • configure_head –

    Build the query-based TransFusion detection head.

  • decode –

    Decode a forward output into raw candidate detections (see TransFusionHead.decode).

  • 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 [source] ΒΆ

configure_vfe() -> DynamicMeanVFE

Build the dynamic mean voxel feature encoder.

configure_backbone_3d [source] ΒΆ

configure_backbone_3d() -> LION3DBackbone

Build the sparse 3D linear group RNN backbone.

configure_backbone [source] ΒΆ

configure_backbone() -> BEVResidualBackbone

Build the residual 2D BEV backbone.

configure_head [source] ΒΆ

configure_head() -> TransFusionHead

Build the query-based TransFusion detection head.

decode [source] ΒΆ

Decode a forward output into raw candidate detections (see TransFusionHead.decode).

reset_classifier [source] ΒΆ

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.

window_partition [source] ΒΆ

window_partition(
    pos: Tensor,
    sparse_shape: Sequence[int],
    window_shape: Sequence[int],
    shift: bool,
) -> Tuple[Tensor, Tensor, Tensor]

Assign each voxel to a 3D window and return its in-window offset (get_window_coors_shift_v2).

Each voxel is assigned to a 3D window of size window_shape (optionally shifted by half a window, Swin-style), yielding two flat per-voxel window keys: one ordering windows \(x\)-major (win_index_x) and one \(y\)-major (win_index_y). The within-window \((z, y, x)\) offset is returned for the intra-window order.

Parameters:

  • pos (Tensor) –

    Voxel indices \((N, 4)\) as \((\text{batch}, z, y, x)\).

  • sparse_shape (Sequence[int]) –

    Grid extent \((z, y, x)\).

  • window_shape (Sequence[int]) –

    Window size \((w_x, w_y, w_z)\).

  • shift (bool) –

    Whether to offset windows by half their size.

Returns:

  • Tuple[Tensor, Tensor, Tensor] –

    (win_index_x, win_index_y, offsets) of shapes \((N,)\), \((N,)\), \((N, 3)\).