lion
LION detection model.

Classes:
-
BiMambaβBidirectional Mamba mixer used by the LION linear group RNN (
Mambain 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
TransformerDecoderLayerover \((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] ΒΆ
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] ΒΆ
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] ΒΆ
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;
-1keepsdim. -
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
diffusionis set.
PatchExpanding3D [source] ΒΆ
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
LIONBlockstages. -
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_classesuse \(> 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
predictand each gets its own circular NMS task indecode. 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).
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:
-
TransFusionHeadOutputβThe raw per-query predictions and the dense heatmap.
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:
-
preds_dicts(TransFusionHeadOutput) βThe dict returned by
forward/predict.
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_classesoutputs.
Attributes:
-
num_features(int) βChannel count \(C\) of the BEV feature map entering the head.
num_features property ΒΆ
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(out: TransFusionHeadOutput) -> Detection3D
Decode a forward output into raw candidate detections (see TransFusionHead.decode).
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)\).