transfusion

TransFusion detection loss: Hungarian-matched query targets and a dense center-heatmap objective.

Classes:

TransFusionLoss –

Query-based TransFusion detection loss (dense heatmap, matched classification, box, IoU rescore).

TransFusionLoss ¶

TransFusionLoss(
    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,
    hungarian_cls_cost: float = 0.15,
    hungarian_reg_cost: float = 0.25,
    hungarian_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,
    bbox_weight: float = 0.25,
    hm_weight: float = 1.0,
    iou_weight: float = 0.5,
    focal_alpha: float = 0.25,
    focal_gamma: float = 2.0,
)

Bases: Module

Query-based TransFusion detection loss (dense heatmap, matched classification, box, IoU rescore).

Reference: 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.

Note

nuScenes ground-truth boxes carry no velocity (\((K, 7)\)), so the velocity targets are zero. The default code_weights zeroes the last two (velocity) codes to leave that branch unsupervised.

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.
hungarian_cls_cost (float, default: 0.15 ) –

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

Weight of the normalized center-\(L_1\) term in the matching cost.
hungarian_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.
bbox_weight (float, default: 0.25 ) –

Multiplier on the box-regression term.
hm_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 components.

forward ¶

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

Compute the TransFusion loss and its components.

Parameters:

output (Dict[str, Tensor]) –

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)\).
batch (Dict[str, Any]) –

Packed ground truth (DataKeys.BOX \((K, 7)\) full-extent, DataKeys.LABEL \((K,)\) \(0\)-based, DataKeys.BATCH_BOX \((K,)\) per-box scene index).

Returns:

Dict[str, Tensor] –

A dict with the scalar loss (to backprop) and detached heatmap_loss, cls_loss,
Dict[str, Tensor] –

bbox_loss, iou_loss diagnostics.