transfusion

Loss of the query-based TransFusionHead (LION): Hungarian-matched query targets and a dense center heatmap.

Classes:

TransFusionHeadLoss –

Loss of the TransFusionHead (dense heatmap, matched classification, box, IoU rescore).

TransFusionHeadLoss [source] ¶

TransFusionHeadLoss(
    num_classes: int,
    *,
    point_cloud_range: Sequence[float],
    voxel_size: Sequence[float],
    feature_map_stride: int,
    num_proposals: int = 200,
    gaussian_overlap: float = 0.1,
    min_radius: int = 2,
    matcher_cls_cost: float = 0.15,
    matcher_reg_cost: float = 0.25,
    matcher_iou_cost: float = 0.25,
    code_weights: Sequence[float] = (
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        0.0,
        0.0,
    ),
    cls_weight: float = 1.0,
    loc_weight: float = 0.25,
    heatmap_weight: float = 1.0,
    iou_weight: float = 0.5,
    focal_alpha: float = 0.25,
    focal_gamma: float = 2.0,
)

Bases: Module

Loss of the TransFusionHead (dense heatmap, matched classification, box, IoU rescore).

Reference: TransFusion: Robust LiDAR-Camera Fusion for 3D Object Detection with Transformers (Bai et al., 2022).

The head predicts a dense per-class BEV heatmap plus a fixed set of object queries, each carrying a class logit vector and a box code. Four terms are summed:

Heatmap: the ground-truth box centers are splatted onto a per-class BEV Gaussian map and the dense heatmap is supervised by the penalty-reduced center focal loss.
Classification: every scene's queries are decoded to boxes and matched to the ground truth by a per-scene Hungarian assignment (cost: focal classification + normalized center \(L_1\) + 3D IoU). The per-query class logits are then trained by sigmoid focal loss over one-hot targets (background for unmatched queries), normalized by the positive count.
Box regression: code-weighted \(L_1\) over the \(10\)-dim box code \((x, y, z + d_z / 2, \log d_x, \log d_y, \log d_z, \sin\theta, \cos\theta, v_x, v_y)\) at the matched queries.
IoU rescore: an \(L_1\) term regressing the per-query iou branch toward \(2 \cdot \text{IoU}_{3D} - 1\) between each matched query's decoded box and its ground-truth box.

The loss holds no reference to the model: the grid geometry is rebuilt from the constructor params.

Parameters:

num_classes (int) –

Number of foreground classes (heatmap channels and query logits).
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) –

Stride from the voxel grid to the BEV feature map.
num_proposals (int, default: 200 ) –

Number of object queries per scene.
gaussian_overlap (float, default: 0.1 ) –

Min-overlap passed to the Gaussian-radius solver.
min_radius (int, default: 2 ) –

Lower clamp on the integer splat radius.
matcher_cls_cost (float, default: 0.15 ) –

Weight of the focal classification term in the matching cost.
matcher_reg_cost (float, default: 0.25 ) –

Weight of the normalized center-\(L_1\) term in the matching cost.
matcher_iou_cost (float, default: 0.25 ) –

Weight of the 3D-IoU term in the matching cost.
code_weights (Sequence[float], default: (1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0) ) –

Per-code regression weight, length \(10\); the last two (velocity) default to \(0\).
cls_weight (float, default: 1.0 ) –

Multiplier on the classification term.
loc_weight (float, default: 0.25 ) –

Multiplier on the box-regression term.
heatmap_weight (float, default: 1.0 ) –

Multiplier on the heatmap term.
iou_weight (float, default: 0.5 ) –

Multiplier on the IoU-rescore term.
focal_alpha (float, default: 0.25 ) –

Focal positive/negative balance (classification loss and matching cost).
focal_gamma (float, default: 2.0 ) –

Focal focusing exponent (classification loss and matching cost).

Methods:

forward –

Compute the TransFusion loss and its terms.

forward [source] ¶

forward(
    output: TransFusionHeadOutput, data: Dict[str, Any]
) -> Dict[str, Tensor]

Compute the TransFusion loss and its terms.

Parameters:

output (TransFusionHeadOutput) –

The head's raw output: per-query center \((B, 2, Q)\), height \((B, 1, Q)\), dim \((B, 3, Q)\), rot \((B, 2, Q)\), vel \((B, 2, Q)\), iou \((B, 1, Q)\) and heatmap \((B, C, Q)\) class logits, plus the dense dense_heatmap \((B, C, H, W)\).
data (Dict[str, Any]) –

Packed ground truth (DataKeys.BOX \((K, 7)\) full-extent, or \((K, 9)\) with the velocity in the last two columns, DataKeys.LABEL \((K,)\) \(0\)-based, DataKeys.BATCH_BOX \((K,)\) per-box scene index).

Returns:

Dict[str, Tensor] –

A dict with the scalar loss and detached heatmap_loss, cls_loss, box_loss, iou_loss.