用DETR实现端到端目标检测：从原理到自定义数据集实战
当目标检测遇上Transformer，传统方法中的Anchor设计突然显得如此多余。DETR（Detection Transformer）的出现，不仅简化了整个检测流程，更让我们看到了视觉任务中注意力机制的强大潜力。本文将带你深入理解这一创新架构，并手把手完成从数据准备到模型部署的全过程。
1. 为什么DETR是目标检测的革命者？
传统目标检测方法如Faster R-CNN和YOLO都依赖复杂的Anchor设计，这些预定义的边界框不仅需要精心调整超参数，还引入了大量计算开销。DETR的核心突破在于：
- 完全摒弃Anchor机制：使用Transformer的注意力机制直接预测目标位置
- 真正的端到端训练：无需后处理NMS（非极大值抑制）
- 统一的任务处理方式：将检测视为集合预测问题
下表对比了传统方法与DETR的关键差异：
| 特性 | Anchor-based方法 | DETR |
|---|---|---|
| 需要预定义Anchor | 是 | 否 |
| 需要NMS后处理 | 是 | 否 |
| 训练流程 | 多阶段 | 端到端 |
| 对小目标检测效果 | 一般 | 中等 |
| 计算复杂度 | 中等 | 较高 |
注意：虽然DETR的计算开销较大，但其简洁的架构使其在多项基准测试中达到了与Faster R-CNN相当的精度。
2. 准备自定义COCO格式数据集
要让DETR识别你的专属目标类别，首先需要构建符合COCO格式的训练数据。以下是关键步骤：
图像采集与标注
- 确保每张图像至少包含一个目标实例
- 使用LabelImg或CVAT等工具进行边界框标注
- 保存为JSON格式的标注文件
构建COCO数据结构
{ "images": [{ "id": int, "width": int, "height": int, "file_name": str }], "annotations": [{ "id": int, "image_id": int, "category_id": int, "bbox": [x,y,width,height], "area": float, "iscrowd": 0 }], "categories": [{ "id": int, "name": str }] }- 数据集目录结构
custom_coco/ ├── annotations/ │ ├── instances_train.json │ └── instances_val.json └── images/ ├── train/ └── val/3. 模型配置与权重调整
DETR官方提供了基于ResNet-50的预训练模型，但我们需要针对自定义类别进行调整：
下载预训练权重
- 官方模型通常使用COCO的91个类别
- 需修改分类头以适应你的类别数量
调整类别数的关键代码
import torch # 加载预训练权重 pretrained_weights = torch.load('detr-r50.pth') # 修改为你的类别数（示例为3类） num_classes = 3 pretrained_weights["model"]["class_embed.weight"].resize_(num_classes+1, 256) pretrained_weights["model"]["class_embed.bias"].resize_(num_classes+1) # 保存调整后的权重 torch.save(pretrained_weights, "detr_r50_%d.pth"%num_classes)- 修改模型配置文件
- 在
models/detr.py中更新num_classes参数 - 调整
main.py中的训练超参数：
- 在
parser.add_argument('--lr', default=1e-4, type=float) parser.add_argument('--batch_size', default=4, type=int) parser.add_argument('--epochs', default=300, type=int) parser.add_argument('--num_queries', default=100, type=int)4. 训练与推理实战
准备好数据和模型后，就可以开始训练你的自定义检测器了：
- 启动训练
python main.py \ --dataset_file "coco" \ --coco_path "/path/to/your/custom_coco" \ --output_dir "output" \ --resume "detr_r50_3.pth" \ --num_classes 3监控训练过程
- 使用TensorBoard记录损失曲线
- 关注分类损失和边界框回归损失的平衡
- 典型训练需要300个epoch左右
推理代码示例
def detect(image_path, model, transform): img = Image.open(image_path) img_tensor = transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs = model(img_tensor) # 处理输出结果 probas = outputs['pred_logits'].softmax(-1)[0, :, :-1] keep = probas.max(-1).values > 0.7 # 置信度阈值 # 转换边界框格式 bboxes_scaled = rescale_bboxes( outputs['pred_boxes'][0, keep].cpu(), img.size ) return probas[keep], bboxes_scaled5. 性能优化与实用技巧
在实际应用中，我们发现几个提升DETR效果的关键点：
- 学习率策略：使用带warmup的余弦退火调度
- 数据增强：适当增加随机裁剪和颜色扰动
- 查询数量：根据目标密集程度调整num_queries
- 注意力可视化：理解模型关注区域
对于小目标检测的局限性，可以考虑：
- 增加高分辨率特征图
- 使用Deformable DETR等改进版本
- 在损失函数中增加对小目标的权重
训练完成后，模型可以直接部署到生产环境，无需复杂的后处理流程。这种端到端的特性使得DETR在实际应用中展现出独特的优势，特别是在需要快速迭代新类别的场景中。
