Skip to content

3DETR

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

Classes:

  • DETR3DLoss –

    3DETR Hungarian set-prediction detection loss.

DETR3DLoss 

DETR3DLoss(
    num_classes: int,
    num_angle_bin: 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,
    loss_giou_weight: float = 0.0,
    loss_sem_cls_weight: float = 1.0,
    loss_no_object_weight: float = 0.2,
    loss_angle_cls_weight: float = 0.1,
    loss_angle_reg_weight: float = 0.5,
    loss_center_weight: float = 5.0,
    loss_size_weight: float = 1.0,
)

Bases: Module

3DETR Hungarian set-prediction detection loss.

Reference: 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 loss_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_angle_bin (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.

  • loss_giou_weight (float, default: 0.0 ) –

    Weight of the GIoU term in the total. The reference trains with \(0\) (the GIoU drives only the matcher); note the rotated-box GIoU (scenes with non-zero headings) is computed without gradients, so a non-zero weight trains only axis-aligned scenes.

  • loss_sem_cls_weight (float, default: 1.0 ) –

    Weight of the semantic-classification term in the total.

  • loss_no_object_weight (float, default: 0.2 ) –

    Cross-entropy weight of the background class.

  • loss_angle_cls_weight (float, default: 0.1 ) –

    Weight of the heading-bin classification term in the total.

  • loss_angle_reg_weight (float, default: 0.5 ) –

    Weight of the heading-residual regression term in the total.

  • loss_center_weight (float, default: 5.0 ) –

    Weight of the center term in the total.

  • loss_size_weight (float, default: 1.0 ) –

    Weight of the size term in the total.

Methods:

  • forward –

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

forward 

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

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

Parameters:

  • output (Dict[str, Any]) –

    A training-mode DETR3DTrainOutput: 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.

  • batch (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 loss_sem_cls,

  • Dict[str, Tensor] –

    loss_center, loss_size, loss_angle_cls, loss_angle_reg, loss_giou, loss_cardinality

  • Dict[str, Tensor] –

    (each summed over decoder layers).