threedetr

3DETR set-prediction detection loss: Hungarian query-to-object matching with per-layer aux losses.

Classes:

ThreeDETRLoss –

3DETR Hungarian set-prediction detection loss.

ThreeDETRLoss [source] ¶

ThreeDETRLoss(
    num_classes: int,
    *,
    num_heading_bins: int,
    matcher_cls_cost: float = 1.0,
    matcher_giou_cost: float = 2.0,
    matcher_center_cost: float = 0.0,
    matcher_objectness_cost: float = 0.0,
    giou_weight: float = 0.0,
    sem_cls_weight: float = 1.0,
    no_object_weight: float = 0.2,
    angle_cls_weight: float = 0.1,
    angle_reg_weight: float = 0.5,
    center_weight: float = 5.0,
    size_weight: float = 1.0,
)

Bases: Module

3DETR Hungarian set-prediction detection loss.

Reference: An End-to-End Transformer Model for 3D Object Detection (Misra et al., 2021).

Object queries are matched to ground-truth boxes one-to-one per scene by a Hungarian assignment whose cost combines the negative predicted class probability, the negative generalized 3D IoU, the \(L_1\) center distance (in the per-scene min-max normalized frame) and the negative objectness. Every decoder layer is supervised (the last layer plus the intermediate layers as auxiliary outputs), each with the same weighted objective, and the per-layer losses are summed:

Semantic classification: per-query weighted cross-entropy over the \(C + 1\) class logits, with unmatched queries assigned the background slot and that slot down-weighted by no_object_weight.
Center: \(L_1\) distance between matched query and box centers in the normalized frame.
Size: \(L_1\) distance between matched query and box sizes in the normalized frame.
Angle: cross-entropy on the heading bin plus a Huber loss on the in-bin residual, over matches.
GIoU: \(1 - \text{gIoU}_{3D}\) between matched query and box, over matches.
Cardinality: the \(L_1\) error between the count of non-background queries and the object count (logged only, never optimized).

Ground truth is read packed from the batch (full-extent \((K, 7)\) boxes with counter-clockwise headings, plus per-box classes) and densified per scene; the headings are negated into the model's native heading space before binning. The normalization uses the model's point_cloud_dims, so the loss holds no reference to the model.

Parameters:

num_classes (int) –

Number of semantic classes (the class head predicts one extra background slot).
num_heading_bins (int) –

Heading-angle bins (\(1\) for axis-aligned ScanNet, \(12\) for oriented SUN RGB-D).
matcher_cls_cost (float, default: 1.0 ) –

Matcher weight on the negative class probability.
matcher_giou_cost (float, default: 2.0 ) –

Matcher weight on the negative generalized 3D IoU.
matcher_center_cost (float, default: 0.0 ) –

Matcher weight on the normalized-center \(L_1\) distance.
matcher_objectness_cost (float, default: 0.0 ) –

Matcher weight on the negative objectness.
giou_weight (float, default: 0.0 ) –

Weight of the GIoU term in the total; \(0\) (the published recipe) keeps the GIoU in the matcher only. The rotated-box GIoU (scenes with non-zero headings) is computed without gradients, so a non-zero weight trains only axis-aligned scenes. Its values differ slightly from implementations that skip the polygon intersection of box pairs whose corner-spanned axis-aligned rectangles do not overlap: box3d_overlap intersects every pair exactly.
sem_cls_weight (float, default: 1.0 ) –

Weight of the semantic-classification term in the total.
no_object_weight (float, default: 0.2 ) –

Cross-entropy weight of the background class.
angle_cls_weight (float, default: 0.1 ) –

Weight of the heading-bin classification term in the total.
angle_reg_weight (float, default: 0.5 ) –

Weight of the heading-residual regression term in the total.
center_weight (float, default: 5.0 ) –

Weight of the center term in the total.
size_weight (float, default: 1.0 ) –

Weight of the size term in the total.

Methods:

forward –

Compute the 3DETR set-prediction loss and its terms.

forward [source] ¶

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

Compute the 3DETR set-prediction loss and its terms.

Parameters:

output (ThreeDETRTrainOutput) –

A training-mode ThreeDETRTrainOutput: aux_outputs (a per-decoder-layer list of head dicts with sem_cls_logits, sem_cls_prob, objectness_prob, center_normalized, center_unnormalized, size_normalized, size_unnormalized, angle_logits, angle_residual_normalized, angle_continuous) and point_cloud_dims.
data (Dict[str, Any]) –

Packed ground truth: DataKeys.BOX \((K, 7)\) full-extent boxes with counter-clockwise headings, DataKeys.LABEL \((K,)\) per-box classes and DataKeys.BATCH_BOX \((K,)\) per-box scene index.

Returns:

Dict[str, Tensor] –

A dict with the scalar loss (summed over decoder layers) and detached sem_cls_loss,
Dict[str, Tensor] –

center_loss, size_loss, angle_cls_loss, angle_reg_loss, giou_loss and the
Dict[str, Tensor] –

cardinality_error diagnostic (each summed over decoder layers).