CVPR2025 DEIM 目标检测模型训练教程

一、下载 CVPR2025 DEIM 官方源码
DEIM 项目源码地址:
[CVPR 2025] DEIM: DETR with Improved Matching for Fast Convergence (GitHub)

二、配置虚拟环境
第一步:创建虚拟环境
conda create -n deim python=3.10
第二步:激活虚拟环境
conda activate deim
第三步:安装 PyTorch
建议不要安装太新版本。
PyTorch 官网:Previous PyTorch Versions
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia

第四步:安装依赖包
复制以下命令到控制台终端运行:
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install torchsummary -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnx==1.14.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxruntime==1.15.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install pycocotools==2.0.7 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install faster-coco-eval==1.6.5 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install PyYAML -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install tensorboard -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install scipy -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install calflops -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxsim==0.4.36 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxruntime-gpu==1.18.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install timm==1.0.7 thop efficientnet_pytorch==0.7.1 einops grad-cam==1.4.8 dill==0.3.6 albumentations pytorch_wavelets==1.3.0 PyWavelets -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -U openmim -i https://pypi.tuna.tsinghua.edu.cn/simple
mim install mmengine -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install mmcv==2.2.0 -f https://download.openmmlab.com/mmcv/dist/cu121/torch2.3/index.html

三、准备数据集和配置步骤
如果没有数据集,可以去公共数据集网站练习:Roboflow Datasets

3.1 数据集示例
本文以训练 Visdrone2019 无人机数据集为例。
VisDrone 数据集由天津大学等团队开源,包含 12 个类别。本项目忽略'忽略区域'和'others',只使用其余 10 个有效类别。
Visdrone2019-COCO 格式数据集下载链接: https://aistudio.baidu.com/datasetdetail/226107/0
Visdrone2019-YOLO 格式数据集下载链接: https://aistudio.baidu.com/datasetdetail/295374
3.2 放置数据集
将数据集放到 datasets 文件夹里。

3.3 配置数据步骤
第一步: 在 DEIM/configs/dataset 文件中创建 visdrone_detection.yml 文件。

第二步: 修改 visdrone_detection.yml 文件。

完整代码:
task: detection
evaluator:
type: CocoEvaluator
iou_types: ['bbox', ]
num_classes: 10 # your dataset classes
remap_mscoco_category: False
train_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: C:/software/mydemo/Demo/DEIM/datasets/visdrone/train2017
ann_file: C:/software/mydemo/Demo/DEIM/datasets/visdrone/annotations/instances_train2017.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: True
num_workers: 4
drop_last: True
collate_fn:
type: BatchImageCollateFunction
val_dataloader:
type: DataLoader
dataset:
type: CocoDetection
img_folder: C:/software/mydemo/Demo/DEIM/datasets/visdrone/val2017
ann_file: C:/software/mydemo/Demo/DEIM/datasets/visdrone/annotations/instances_val2017.json
return_masks: False
transforms:
type: Compose
ops: ~
shuffle: False
num_workers: 4
drop_last: False
collate_fn:
type: BatchImageCollateFunction
第三步: 修改 DEIM\configs\deim_dfine\dfine_hgnetv2_n_coco.yml 文件。

只需修改其中一行代码:
__include__: [
'../dataset/visdrone_detection.yml',
'../runtime.yml',
'../base/dataloader.yml',
'../base/optimizer.yml',
'../base/dfine_hgnetv2.yml',
]

第四步: 修改 deim_hgnetv2_n_coco.yml 文件。
当使用不同数据集训练时,保存结果文件夹需要重新创建,否则第二次训练后的权重文件会覆盖第一次的。

第五步: 修改 batch_size 大小,根据设备显存灵活设置。
修改 configs/base/dataloader.yml 中的 total_batch_size。

同时修改 DEIM\configs\deim_dfine\dfine_hgnetv2_n_coco.yml 中的 total_batch_size。

第六步: 修改 DEIM\configs\runtime.yml 文件中的参数,方便观察。
checkpoint_freq: 1

到此为止,配置自己的数据集操作已经完成!
四、使用 DEIM 训练数据集
4.1 修改 train.py 脚本参数

4.2 train.py 脚本完整代码
""" DEIM: DETR with Improved Matching for Fast Convergence
Copyright (c) 2024 The DEIM Authors. All Rights Reserved.
---------------------------------------------------------------------------------
Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright (c) 2023 lyuwenyu. All Rights Reserved.
"""
import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import warnings
warnings.filterwarnings('ignore')
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
import argparse
from engine.misc import dist_utils
from engine.core import YAMLConfig, yaml_utils
from engine.solver import TASKS
debug = False
if debug:
import torch
def custom_repr(self):
return f'{{Tensor:{tuple(self.shape)}}} {original_repr(self)}'
original_repr = torch.Tensor.__repr__
torch.Tensor.__repr__ = custom_repr
def main(args, ) -> None:
"""main"""
dist_utils.setup_distributed(args.print_rank, args.print_method, seed=args.seed)
assert not all([args.tuning, args.resume]), \
'Only support from_scrach or resume or tuning at one time'
update_dict = yaml_utils.parse_cli(args.update)
update_dict.update({k: v for k, v in args.__dict__.items() \
if k not in ['update', ] and v is not None})
cfg = YAMLConfig(args.config, **update_dict)
if args.resume or args.tuning:
if 'HGNetv2' in cfg.yaml_cfg:
cfg.yaml_cfg['HGNetv2']['pretrained'] = False
print('cfg: ', cfg.__dict__)
solver = TASKS[cfg.yaml_cfg['task']](cfg)
if args.test_only:
solver.val()
else:
solver.fit()
dist_utils.cleanup()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', default=r'C:\software\mydemo\Demo\DEIM\configs\deim_dfine\deim_hgnetv2_n_coco.yml', type=str)
parser.add_argument('-r', '--resume', type=str, help='resume from checkpoint')
parser.add_argument('-t', '--tuning', type=str, help='tuning from checkpoint')
parser.add_argument('-d', '--device', type=str, help='device')
parser.add_argument('--seed', default=0, type=int, help='exp reproducibility')
parser.add_argument('--use-amp', action='store_true', help='auto mixed precision training')
parser.add_argument('--output-dir', type=str, help='output directoy')
parser.add_argument('--summary-dir', type=str, help='tensorboard summry')
parser.add_argument('--test-only', action='store_true', default=False,)
parser.add_argument('-u', '--update', nargs='+', help='update yaml config')
parser.add_argument('--print-method', type=str, default='builtin', help='print method')
parser.add_argument('--print-rank', type=int, default=0, help='print rank id')
parser.add_argument('--local-rank', type=int, help='local rank id')
args = parser.parse_args()
main(args)
4.3 Windows 系统首次训练 Bug 修复
Windows 系统上第一次训练,100% 会出现这个 bug 问题。

解决方法: 修改文件 DEIM\engine\backbone\hgnetv2.py。

except Exception as e:
print(f"{str(e)}")
logging.error(RED + "CRITICAL WARNING: Failed to load pretrained HGNetV2 model" + RESET)
logging.error(GREEN + f"Please manually download from {download_url} to {local_model_dir}" + RESET)
现在可以正常训练使用 DEIM 官方模型了。
第一次运行默认加载预训练权重,所以需要下载。(建议离线下载比较快,或等待自动在线下载)。演示中关闭了加载预训练权重变量。

五、验证模型效果
使用 val.py 脚本。

DEIM 训练部分 epoch,演示使用 val.py 在验证集或测试集上的效果。

六、推理测试模型
使用 predict.py 脚本。


七、生成论文实验热力图
使用 DEIM 生成热力图脚本生成热力图,论文实验部分增加可视化热力图分析,提高论文质量。

八、总结
本文详细介绍了 CVPR2025 DEIM 目标检测模型的完整使用流程。内容包含:1)源码下载与虚拟环境配置;2)Visdrone2019 数据集准备与配置文件修改;3)训练参数设置及 Windows 系统常见问题解决方案;4)模型验证、推理测试及热力图生成方法。特别针对 Windows 用户提供了本地虚拟环境配置,并详细说明了数据集路径设置、参数调整等关键步骤。通过本文的指导,可以快速掌握 DEIM 模型的使用方法,用于学术论文研究和实验分析。


