DN-DETR: 通过引入查询去噪加速 DETR 训练

作者：Feng Li*、Hao Zhang*、Shilong Liu、Jian Guo、Lionel M.Ni 和 Lei Zhang。

本仓库是 DN-DETR 的官方实现。该工作已被 CVPR 2022 接收（评分 112，口头报告）。代码现已公开。 [CVPR 论文链接] [扩展版论文链接] [中文解读]

最新消息

[2022/12]: 我们在 arXiv 上发布了 DN-DETR 的扩展版本，论文链接为 这里! 我们将去噪训练方法应用于基于 CNN 的 Faster R-CNN 模型、分割模型 Mask2Former，以及 Anchor DETR 和 DETR 等其他 DETR 类模型，以提升这些模型的性能。

[2022/12]: Mask DINO 的代码已开源！Mask DINO 在 COCO 数据集上使用 ResNet-50 和 SwinL 骨干，在不使用额外检测数据的情况下，分别取得了 51.7 和 59.0 的边界框 AP，超越了相同设置下的 DINO！

[2022/11]: 基于 DN-DETR 的 DINO 实现已在本仓库发布。特别感谢 @Vallum! 该优化版本在 ResNet-50 骨干下，仅需 36 个 epoch 即可达到 50.8 ~ 51.0 AP。

[2022/9]: 我们发布了一个工具箱 detrex，其中包含多种最先进的基于 Transformer 的目标检测算法。该工具箱包括性能更优的 DN-DETR。欢迎使用！

[2022/7] DINO 的代码已在此处开放：DINO!

[2022/6]: 我们发布了一种统一的目标检测与分割模型 Mask DINO，在三大分割任务中均取得了最佳成绩（在 COCO 实例分割排行榜 上获得 54.5 AP，在 COCO 全景分割排行榜 上获得 59.4 PQ，在 ADE20K 语义分割排行榜 上获得 60.8 mIoU）! 代码将在此处开放：Mask DINO。

[2022/5]我们的代码现已公开！使用 ResNet-50 骨干，在 COCO 数据集上实现了更好的性能，AP 达到 49.5。

[2022/4] DAB-DETR 的代码已在此处开放：DAB-DETR。

[2022/3]我们创建了一个仓库 Awesome Detection Transformer，用于整理和展示关于目标检测与分割领域 Transformer 相关的论文。欢迎大家关注！

[2022/3]DN-DETR 被选为 CVPR 2022 的 口头报告。

[2022/3]我们还发布了另一项工作 DINO:DETR with Improved DeNoising Anchor Boxes for End-to-End Object Detection，首次将一种 DETR 类模型确立为 COCO 数据集目标检测排行榜 上的 SOTA 模型。该模型同样基于 DN 技术。代码将在此处开放：DINO。

简介

1. 我们提出了一种新颖的去噪训练方法，用于加速 DETR 的训练，并深入理解了 DETR 类方法收敛缓慢的问题。
2. DN 只是一种训练方法，可以轻松集成到多种 DETR 类模型，甚至传统模型中，以提升其性能。
3. DN-DETR 使用 ResNet-50 骨干，在分别经过 12 和 50 个训练 epoch 后，分别达到了 43.4 和 48.6 的 AP。与相同设置下的基线模型相比，DN-DETR 仅用 50% 的训练时间就达到了相当的性能。
4. 经过优化后的模型表现更佳。DN-Deformable-DETR 使用 ResNet-50 骨干时，AP 可达 49.5。

模型

我们在 DAB-DETR 的基础上增加了去噪模块，以加速训练收敛。该模块仅引入少量计算开销，并且在推理阶段会被移除。 我们进行了大量实验来验证去噪训练的有效性，例如收敛曲线的对比。更多实验结果请参阅 我们的论文。

模型库

我们提供了以下模型：DAB-DETR、DAB-Deformable-DETR（仅可变形编码器） 和 DAB-Deformable-DETR（有关详细信息，请参阅 DAB-DETR 的 代码 和 论文）。

您也可以参考我们的

[Google Drive 模型库]

[百度网盘模型库]（提取码：niet）。

50轮次设置

	名称	骨干网络	边界框AP	日志/配置/检查点	在我们的论文中的位置
0	DN-DETR-R50	R50	44.41	Google云端硬盘 / 百度网盘	表1
2	DN-DETR-R50-DC5	R50	46.3	Google云端硬盘 / 百度网盘	表1
5	DN-DAB-可变形DETR (仅可变形编码器)3	R50	48.6	Google云端硬盘 / 百度网盘	表3
6	DN-DAB-可变形DETR-R50-v24	R50	49.5（24轮次时为48.4）	Google云端硬盘 / 百度网盘	优化后的实现，在编码器和解码器中均使用可变形注意力。更多详情请参见DAB-DETR。

12轮次设置

	名称	骨干网络	边界框AP	日志/配置/检查点	在我们的论文中的位置
1	DN-DAB-DETR-R50-DC5(3 pat)2	R50	41.7	Google云端硬盘 / 百度网盘	表2
4	DN-DAB-DETR-R101-DC5(3 pat)2	R101	42.8	Google云端硬盘 / 百度网盘	表2
5	DN-DAB-可变形DETR (仅可变形编码器)3	R50	43.4	Google云端硬盘 / 百度网盘	表2
5	DN-DAB-可变形DETR (仅可变形编码器)3	R101	44.1	Google云端硬盘 / 百度网盘	表2

注释：

• ¹：相较于我们论文中报告的结果（从44.1提升至44.4），这一结果有所提高，因为我们对代码进行了优化。我们并未重新运行其他模型，因此 预计您将获得比我们论文中报告的更好的性能。
• ²：带有标记（3 pat）的模型使用了多模式嵌入进行训练（详情请参考Anchor DETR或DAB-DETR）。
• ³：该模型基于DAB-可变形DETR（仅可变形编码器），这是DAB-DETR的多尺度版本。由于仅在编码器中使用可变形注意力，因此需要16张GPU卡进行训练。
• ⁴：该模型基于DAB-可变形DETR，这是一个使用可变形DETR的优化实现。更多详情请参见DAB-DETR。 建议您使用这个可变形版本，因为它在编码器和解码器中都使用可变形注意力，更加轻量级（即可以用4/8张A100 GPU进行训练），并且收敛速度更快（即在24轮次内达到48.4，与50轮次的DAB-可变形DETR相当）。

使用方法

如何在您自己的模型中使用去噪训练

我们的代码主要遵循DAB-DETR，并添加了用于去噪训练的额外组件，这些组件被封装在一个文件dn_components.py中。主要有3个函数，包括prepare_for_dn、dn_post_proces（前两个函数用于您的检测前向传播中处理去噪部分），以及compute_dn_loss（用于计算去噪损失）。您可以导入这些函数并将其添加到您自己的检测模型中。 如果您希望在自己的检测模型中使用这些功能，也可以对比DN-DETR和DAB-DETR，了解这些函数是如何被加入的。

我们也鼓励您将其应用于其他类似DETR的模型，甚至传统的检测模型，并在此仓库中更新结果。

安装

我们以DAB-DETR项目作为我们的代码库，因此对于DN-DETR无需额外依赖。而对于DN-可变形DETR，您需要手动编译可变形注意力算子。

我们是在python=3.7.3, pytorch=1.9.0, cuda=11.1的环境下测试了我们的模型。其他版本也可能适用。

1. 克隆本仓库

git clone https://github.com/IDEA-Research/DN-DETR.git
cd DN-DETR

2. 安装PyTorch和 torchvision

请按照https://pytorch.org/get-started/locally/上的说明进行操作。

# 示例：
conda install -c pytorch pytorch torchvision

3. 安装其他所需软件包

pip install -r requirements.txt

4. 编译CUDA算子

cd models/dn_dab_deformable_detr/ops
python setup.py build install

# 单元测试（应看到所有检查均为 True）
python test.py
cd ../../..

数据

请下载 COCO 2017 数据集，并按以下方式组织：

COCODIR/
  ├── train2017/
  ├── val2017/
  └── annotations/
  	├── instances_train2017.json
  	└── instances_val2017.json

运行

我们以标准的 DN-DETR-R50 和 DN-Deformable-DETR-R50 为例进行训练和评估。

评估我们的预训练模型

从 此链接 下载我们的 DN-DETR-R50 模型检查点，并执行以下命令。您预计最终的 AP 约为 44.4。

对于我们的 DN-DAB-Deformable-DETR_Deformable_Encoder_Only（在此下载），预期的最终 AP 为 48.6。

对于我们的 DN-DAB-Deformable-DETR（在此下载），预期的最终 AP 为 49.5。

# 对于 dn_detr：44.1 AP；优化后的结果为 44.4 AP
python main.py -m dn_dab_detr \
  --output_dir logs/dn_DABDETR/R50 \
  --batch_size 1 \
  --coco_path /path/to/your/COCODIR \ # 将参数替换为您的 COCO 数据集路径
  --resume /path/to/our/checkpoint \ # 将参数替换为您的检查点路径
  --use_dn \
  --eval

# 对于 dn_deformable_detr：49.5 AP
python main.py -m dn_deformable_detr \
  --output_dir logs/dab_deformable_detr/R50 \
  --batch_size 1 \
  --coco_path /path/to/your/COCODIR \ # 将参数替换为您的 COCO 数据集路径
  --resume /path/to/our/checkpoint \ # 将参数替换为您的检查点路径
  --transformer_activation relu \
  --use_dn \
  --eval
  
# 对于 dn_deformable_detr_deformable_encoder_only：48.6 AP
python main.py -m dn_dab_deformable_detr_deformable_encoder_only 
  --output_dir logs/dab_deformable_detr/R50 \
  --batch_size 1 \
  --coco_path /path/to/your/COCODIR \ # 将参数替换为您的 COCO 数据集路径
  --resume /path/to/our/checkpoint \ # 将参数替换为您的检查点路径
  --transformer_activation relu \
  --num_patterns 3 \  # 使用 3 个模式嵌入
  --use_dn  \
  --eval

训练您自己的模型

同样，您也可以在单进程上训练我们的模型：

# 对于 dn_detr
python main.py -m dn_dab_detr \
  --output_dir logs/dn_DABDETR/R50 \
  --batch_size 1 \
  --epochs 50 \
  --lr_drop 40 \
  --coco_path /path/to/your/COCODIR  # 将参数替换为您的 COCO 数据集路径
  --use_dn

分布式运行

然而，由于训练耗时较长，我们建议在多设备上进行训练。

如果您计划在使用 Slurm 的集群上训练模型，以下是训练示例命令：

# 对于 dn_detr：44.4 AP
python run_with_submitit.py \
  --timeout 3000 \
  --job_name DNDETR \
  --coco_path /path/to/your/COCODIR \
  -m dn_dab_detr \
  --job_dir logs/dn_DABDETR/R50_%j \
  --batch_size 2 \
  --ngpus 8 \
  --nodes 1 \
  --epochs 50 \
  --lr_drop 40 \
  --use_dn

# 对于 dn_dab_deformable_detr：49.5 AP
python run_with_submitit.py \
  --timeout 3000 \
  --job_name dn_dab_deformable_detr \
  --coco_path /path/to/your/COCODIR \
  -m dab_deformable_detr \
  --transformer_activation relu \
  --job_dir logs/dn_dab_deformable_detr/R50_%j \
  --batch_size 2 \
  --ngpus 8 \
  --nodes 1 \
  --epochs 50 \
  --lr_drop 40 \
  --use_dn

# 对于 dn_dab_deformable_detr_deformable_encoder_only：48.6 AP
python run_with_submitit.py \
  --timeout 3000 \
  --job_name dn_dab_deformable_detr_deformable_encoder_only \
  --coco_path /path/to/your/COCODIR \
  -m dn_dab_deformable_detr_deformable_encoder_only \
  --transformer_activation relu \
  --job_dir logs/dn_dab_deformable_detr/R50_%j \
  --num_patterns 3 \ 
  --batch_size 1 \
  --ngpus 8 \
  --nodes 2 \
  --epochs 50 \
  --lr_drop 40 \
  --use_dn

如果您想训练我们的 DC 版本或多模式版本，请添加：

--dilation  # 对于 DC 版本

--num_patterns 3  # 对于 3 种模式

不过，这需要额外的训练资源和内存，即使用 16 个 GPU。

最终的 AP 应该与我们的结果相似或更好，因为我们的优化结果优于论文中报告的性能（例如，我们报告 DN-DETR 的 AP 为 44.1，但新的结果可以达到 44.4）。如果您得到更好的结果也不必惊讶！

我们的训练设置与 DAB-DETR 相同，只是增加了一个参数 --use_dn。您也可以参考 DAB-DETR。

注意：

• 结果对批次大小较为敏感。我们默认使用 16（每张 GPU 2 张图像 × 8 张 GPU）。

或者在单节点上使用多进程运行：

# 对于 dn_dab_detr：44.4 AP
python -m torch.distributed.launch --nproc_per_node=8 \
  main.py -m dn_dab_detr \
  --output_dir logs/dn_DABDETR/R50 \
  --batch_size 2 \
  --epochs 50 \
  --lr_drop 40 \
  --coco_path /path/to/your/COCODIR \
  --use_dn

# 对于 dn_deformable_detr：49.5 AP
python -m torch.distributed.launch --nproc_per_node=8 \
  main.py -m dn_dab_deformable_detr \
  --output_dir logs/dn_dab_deformable_detr/R50 \
  --batch_size 2 \
  --epochs 50 \
  --lr_drop 40 \
  --transformer_activation relu \
  --coco_path /path/to/your/COCODIR \
  --use_dn

链接

我们的工作基于 DAB-DETR。我们还发布了另一个基于 DN-DETR 和 DAB-DETR 的 SOAT 检测模型 DINO。

• DINO：带有改进去噪锚框的 DETR，用于端到端目标检测。 Hao Zhang*, Feng Li*, Shilong Liu*, Lei Zhang, Hang Su, Jun Zhu, Lionel M. Ni, Heung-Yeung Shum。 arxiv 2022。 [论文] [代码]。

• DAB-DETR：动态锚框是 DETR 更好的查询。 Shilong Liu, Feng Li, Hao Zhang, Xiao Yang, Xianbiao Qi, Hang Su, Jun Zhu, Lei Zhang。 国际表征学习会议 (ICLR) 2022。 [论文] [代码]。

许可证

DN-DETR 采用 Apache 2.0 许可证发布。更多信息请参阅 LICENSE 文件。

版权所有 © IDEA。保留所有权利。

根据 Apache 许可证第 2.0 版（“许可证”）授权。除非符合许可证规定，否则不得使用这些文件。您可以在 http://www.apache.org/licenses/LICENSE-2.0 获取许可证副本。

除非适用法律要求或书面协议另有约定，否则根据许可证分发的软件以“AS IS”的基础提供，不附带任何类型的保证或条件，无论是明示的还是暗示的。有关特定语言的权限和限制，请参阅许可证。

Bibtex

如果您认为我们的工作对您的研究有所帮助，请考虑引用以下 BibTeX 条目。

@inproceedings{li2022dn,
  title={Dn-DETR：通过引入查询去噪加速 DETR 训练},
  author={李峰和张浩和刘世龙和郭健和倪利昂和张磊},
  booktitle={IEEE/CVF 计算机视觉与模式识别会议论文集},
  pages={13619--13627},
  year={2022}
}

DN-DETR 快速上手指南

DN-DETR 是一种通过引入查询去噪（Query DeNoising）机制来加速 DETR 训练的创新方法。该方法仅需在训练阶段添加少量计算，推理时自动移除，可显著缩短收敛时间并提升检测性能（如在 COCO 数据集上仅用 50% 的训练轮次即可达到基线模型相当的效果）。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.7.3 (其他版本可能兼容，但已测试此版本)
• PyTorch: 1.9.0
• CUDA: 11.1
• 硬件:
  • 基础模型：单卡或多卡 GPU 均可。
  • 多尺度变形模型 (Deformable Encoder Only)：建议配备 16 张 GPU。
  • 优化版变形模型 (Encoder & Decoder)：建议使用 4 或 8 张 A100 GPU。

注意：如果您计划使用 DN-Deformable-DETR 系列模型，需要手动编译 deformable attention 算子。基础 DN-DETR 无需额外依赖。

安装步骤

1. 克隆代码仓库

git clone https://github.com/IDEA-Research/DN-DETR.git
cd DN-DETR

2. 安装 PyTorch 和 TorchVision

请访问 PyTorch 官网 获取适合您环境的安装命令。以下是基于 CUDA 11.1 的示例：

pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

3. 安装其他依赖

进入项目目录并安装 requirements：

pip install -r requirements.txt

4. 编译 Deformable Attention 算子 (仅针对变形模型)

如果您需要使用 DN-DAB-Deformable-DETR 模型，必须编译自定义算子：

cd models/DAB_DETR/ops
python setup.py build install
cd ../../../

基本使用

模型推理与测试

本项目基于 DAB-DETR 构建。您可以直接下载预训练权重进行评测。

下载预训练模型：

• Google Drive: 模型库链接
• 百度网盘: 模型库链接 (提取码: niet)

运行评估脚本： 以 ResNet-50 骨干网络的 50 epoch 模型为例：

python main.py \
  --output_dir ./logs/dn_detr_r50 \
  --config_file configs/dn_detr_r50.yaml \
  --eval \
  --resume path/to/checkpoint.pth

(注：具体配置文件名和参数请参考 configs/ 目录下的实际文件)

将去噪训练应用到您的模型

DN-DETR 的核心逻辑封装在 models/DAB_DETR/dn_components.py 文件中。若要在自己的检测模型中集成此功能，主要涉及以下三个函数：

1. prepare_for_dn: 在检测前向传播中准备去噪数据。
2. dn_post_process: 处理去噪部分的输出。
3. compute_dn_loss: 计算去噪损失。

集成示例思路： 您可以在自己的模型 forward 函数中导入并调用这些组件：

from models.DAB_DETR.dn_components import prepare_for_dn, dn_post_process, compute_dn_loss

# 1. 准备去噪输入
dn_meta, known_labels, known_boxes = prepare_for_dn(targets, ...)

# 2. 在模型前向传播中使用 dn_meta 生成带噪声的查询

# 3. 计算损失时加入 dn loss
loss_dict = criterion(outputs, targets)
loss_dn = compute_dn_loss(...)
final_loss = sum(loss_dict.values()) + loss_dn

建议参考本仓库中 DN-DETR 与 DAB-DETR 的代码差异，以了解具体的插入位置和调用方式。

推荐模型选择

• 追求轻量与快速收敛：推荐使用 DN-DAB-Deformable-DETR-R50-v2。该版本在 Encoder 和 Decoder 中均使用了 deformable attention，仅需 24 个 epoch 即可达到 48.4 AP，且显存占用更优（支持 4/8 卡训练）。
• 复现论文基准：可选择表格中对应的 R50 或 R101 backbone 模型。