MiniDETR:简化DETR实现

Posted on In Views:

本项目实现了 DETR（DEtection TRansformer）模型的一个简化版本，并最终发布在了 GitHub 上：MiniDETR 项目地址。

这是一个目标检测模型的完整训练框架。

本篇博客总结了我对 DETR 的理解、模型的核心模块、关键实现细节、训练调优策略，以及一些工程实践经验。

DETR核心思想简述

传统目标检测模型（如 Faster R-CNN、YOLO）通常依赖 anchor box、NMS 等启发式方法。而 DETR 的创新点在于：

  • 目标检测 = 集合预测问题，由匈牙利匹配直接解决框-目标对齐问题
  • 用 Transformer 进行全局建模，输入图像特征后输出一组固定长度的 object queries，代表可能的目标
  • End-to-End 无需手工设计后处理模块

这使得 DETR 的架构极度简洁统一，同时也为 Transformer 在视觉领域开辟了新的道路。

实现结构概览

我实现的 Mini-DETR 模型保留了 DETR 的核心思路，并针对小显存设备（≤8GB）进行了简化与优化：

  • 使用 ResNet50 作为 backbone，输出 feature map
  • 用 1×1 Conv 将 channel 压缩至 hidden_dim=256
  • 加入 位置编码（sinusoidal），并送入自实现的 TransformerEncoder + TransformerDecoder
  • decoder 接收 num_queries=100 个 object query，输出每个目标的分类与 bbox
  • 匹配预测与GT使用 HungarianMatcher
  • 损失包括：CrossEntropy + L1 + GIoU

项目结构设计

工程采用主流 PyTorch 项目结构划分，具备良好的模块化和可扩展性：

  • models/: 模型主干与 Transformer 组件
  • data/: VOC数据加载、自定义 JSON 格式解析、图像增广
  • utils/: 包含 matcher、损失函数、评估工具（mAP 计算）
  • train.py: 包含训练循环、tqdm 进度条、warmup 策略、checkpoints、CSV日志等
  • eval.py: 进行 mAP 评估，兼容 VOC JSON 数据格式

这次我特别注重代码组织清晰与注释完整，每个模块都可以独立理解与运行调试。

实现关键细节分享

1. 匈牙利匹配的理解与实现

DETR的精髓在于将检测任务视为集合预测，使用匈牙利算法做 bipartite matching。 我实现了一个简化版本，匹配代价由分类差异（CE）+ 回归差异（L1）组成，输出预测-目标的一一对应关系。

2. 位置编码（Positional Embedding）

我使用了与原文一致的 正弦-余弦编码，编码二维位置，并添加到特征上。位置编码模块是完全手写的。

3. Transformer模块自实现

4. 支持 warmup 与差分学习率

在训练阶段，我加入了：

  • 学习率 warmup（前1000步线性增长）
  • backbone 冻结 / 解冻机制
  • backbone 与 transformer 使用不同的学习率

5. mAP 评估工具

没有使用 pycocotools，自己实现了 mAP（IoU=0.5）计算逻辑，包括：

  • IoU计算
  • TP/FP判定
  • Precision-Recall 曲线
  • 11-point mAP 计算

成果

在VOC训练集上训练约50 epochs，模型能够初步检测出常见目标，mAP 在 0.4~0.5 左右（IoU=0.5）。 虽然精度不高，但这个结果已经验证了模型的通路是正确的，对于学习和调试是完全可用的。

更重要的是，我从零构建出了一个完整检测框架，这比直接套用库要收获大得多

项目地址

完整代码已开源在 GitHub：

https://github.com/QianQing26/MiniDETR