Fix D-FINE / RT-DETR main loss being computed over the denoising queries

[stefan-it]

What does this PR do?

Fixes a training bug in the D-FINE and RT-DETR (incl. RT-DETRv2) losses: the main, Hungarian-matched loss term was computed over the contrastive-denoising queries as well, which starves the normal queries of the inference layer of positive supervision.

The bug

DFineForObjectDetectionLoss (src/transformers/loss/loss_d_fine.py) and RTDetrForObjectDetectionLoss (src/transformers/loss/loss_rt_detr.py, also mapped for RTDetrV2ForObjectDetection) receive the last decoder layer's logits / pred_boxes from the model. In training mode with num_denoising > 0 these tensors contain the contrastive denoising (CDN) queries followed by the normal queries, e.g. [batch, 200 + 300, num_labels]. Both functions split the CDN queries off for the auxiliary and dn_* terms (torch.split(..., dn_num_split, dim=2)), but never for the main term:

outputs_loss["logits"] = logits            # still [batch, 200 + 300, num_labels] during training
outputs_loss["pred_boxes"] = pred_boxes    # -> Hungarian matching over all 500 queries

The positive CDN queries are initialized from lightly noised ground-truth boxes, so the matcher of the main term assigns most targets to them. The normal queries of the last decoder layer, the only ones used at inference (eval_idx = -1), then get almost no positive signal from the main term. Training losses look healthy while validation mAP stalls with low, badly calibrated scores and class confusion, which makes this hard to spot.

The reference implementations split the denoising queries before building the main term:

• D-FINE: https://github.com/Peterande/D-FINE/blob/master/src/zoo/dfine/dfine_decoder.py (torch.split(out_logits, dn_meta["dn_num_split"], dim=2) → "pred_logits": out_logits[-1])
• RT-DETR: https://github.com/lyuwenyu/RT-DETR/blob/main/rtdetrv2_pytorch/src/zoo/rtdetr/rtdetr_decoder.py

loss_deimv2.py already does this for DEIMv2, so it is not affected.

Measurements

Fine-tuning ustc-community/dfine-nano-coco on a 4-class document layout dataset (6.4k images) with the original D-FINE hyper-parameters. Matching of the main term on a real training batch (8 images, 12 ground-truth boxes) after 7 epochs:

out = model(pixel_values=pixel_values, labels=labels)       # model.train()
criterion = DFineLoss(model.config)
indices = criterion.matcher({"logits": out.logits, "pred_boxes": out.pred_boxes.clamp(0, 1)}, labels)
matched = torch.cat([src for src, _ in indices])
num_dn = out.denoising_meta_values["dn_num_split"][0]        # 200
print((matched < num_dn).sum().item(), "/", len(matched))    # -> 10 / 12 targets matched to denoising queries

The last-layer loss_vfl restricted to the normal queries was 2.55, while the value entering the training loss was 0.92.

Class-aware COCO AP on the validation split (no score threshold), same data, hyper-parameters and schedule:

epoch	before (stock loss)	after (this PR)	original D-FINE repo
3	0.17	0.51	0.14
7	0.25	0.68	0.41
15	–	0.73	0.60

With the stock loss the model localized boxes fine but labelled nearly everything as a single class with low scores; the class-agnostic AP at the 0.3 score threshold plateaued around 0.4 for 45 epochs, while the original D-FINE repo reaches 0.71 on that metric.

The fix

At the top of both loss functions, before the config.auxiliary_loss branch (so it also applies with auxiliary losses disabled):

if denoising_meta_values is not None:
    _, logits = torch.split(logits, denoising_meta_values["dn_num_split"], dim=1)
    _, pred_boxes = torch.split(pred_boxes, denoising_meta_values["dn_num_split"], dim=1)

Reproduction

Fails on main with AssertionError: Scalars are not close!, passes with this PR:

import torch
from transformers import DFineConfig, DFineForObjectDetection
from transformers.loss.loss_d_fine import DFineLoss

config = DFineConfig(num_labels=4, num_denoising=10, num_queries=30)
model = DFineForObjectDetection(config).train()
pixel_values = torch.rand(2, 3, 256, 256)
labels = [{"class_labels": torch.randint(0, 4, (3,)), "boxes": torch.rand(3, 4)} for _ in range(2)]
out = model(pixel_values=pixel_values, labels=labels)
num_dn = out.denoising_meta_values["dn_num_split"][0]
ref = DFineLoss(config)({"logits": out.logits[:, num_dn:], "pred_boxes": out.pred_boxes[:, num_dn:].clamp(0, 1)}, labels)
for k in ("loss_vfl", "loss_bbox", "loss_giou"):
    torch.testing.assert_close(out.loss_dict[k], ref[k])

Tests

test_main_loss_excludes_denoising_queries in tests/models/d_fine/test_modeling_d_fine.py and tests/models/rt_detr/test_modeling_rt_detr.py trains a small model with denoising enabled and asserts that the main loss terms equal the loss recomputed on the normal queries alone. Both fail on main and pass with this PR. The full (non-slow) D-FINE, RT-DETR and RT-DETRv2 test files pass:

pytest tests/models/d_fine/test_modeling_d_fine.py tests/models/rt_detr/test_modeling_rt_detr.py tests/models/rt_detr_v2/test_modeling_rt_detr_v2.py

Side note, not part of this PR: the original D-FINE also applies the fine-grained localization loss (loss_fgl) to the last decoder layer, whereas the HF main term only has loss_vfl / loss_bbox / loss_giou. That is a smaller fidelity gap and can be addressed separately.

The bug was found while porting a D-FINE fine-tuning pipeline to Transformers, with the analysis done by Claude Fable 5.1 (Claude Code).

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a Github issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

@qubvel

🤖 Generated with Claude Code

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: d_fine, rt_detr

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33884396458:1
Result: success | Jobs: 16 | Tests: 185,241 | Failures: 0 | Duration: 16h 12m