Windows10下DETR目标检测实战：从COCO到自定义数据集的完整迁移指南

在计算机视觉领域，目标检测一直是核心任务之一。传统的检测方法如Faster R-CNN、YOLO等虽然效果显著，但往往需要复杂的后处理流程。Facebook AI提出的DETR（DEtection TRansformer）首次将Transformer架构引入目标检测领域，实现了端到端的检测流程。本文将带你在Windows10环境下，完成从COCO标准数据集到自定义数据集的完整迁移过程。

1. 环境准备与依赖安装

在Windows10上搭建DETR开发环境需要特别注意几个关键点。首先确保你的系统已安装Python 3.7+和PyTorch 1.5+。推荐使用Anaconda创建独立环境以避免依赖冲突：

conda create -n detr python=3.8
conda activate detr

安装基础依赖时，最容易出现问题的是pycocotools。Windows环境下建议使用预编译版本：

pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI

完整的依赖列表包括：

• torch>=1.5.0
• torchvision>=0.6.0
• opencv-python
• scipy
• pycocotools

注意：如果遇到VC++编译错误，需要安装Visual Studio 2019的C++构建工具，勾选"使用C++的桌面开发"工作负载。

2. 数据集准备与格式转换

DETR默认使用COCO格式的数据集，但实际项目中我们往往需要处理自定义数据。以下是VOC格式转换为COCO格式的关键步骤：

1. 使用LabelImg等工具标注数据，生成VOC格式的XML文件
2. 创建类别映射文件classes.txt，每行一个类别名
3. 运行转换脚本生成COCO格式的JSON标注文件

典型的目录结构应如下：

custom_dataset/
├── annotations/
│   ├── instances_train2017.json
│   └── instances_val2017.json
└── train2017/
    ├── 000001.jpg
    ├── 000002.jpg
    └── ...

转换脚本的核心逻辑是解析XML并重构为COCO JSON格式。关键数据结构包括：

{
    "images": [{
        "file_name": "000001.jpg",
        "height": 480,
        "width": 640,
        "id": 1
    }],
    "annotations": [{
        "area": 5000,
        "iscrowd": 0,
        "image_id": 1,
        "bbox": [x,y,width,height],
        "category_id": 1,
        "id": 1
    }],
    "categories": [{
        "id": 1,
        "name": "cat"
    }]
}

3. 模型适配与权重调整

DETR预训练权重是基于COCO的91类别设计的。迁移到自定义数据集时，需要调整分类头：

import torch
pretrained_weights = torch.load('detr-r50.pth')

# 修改分类头维度
num_classes = 5  # 自定义类别数+1(背景)
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-custom.pth")

同时需要修改模型配置文件detr.py中的类别数：

# 修改num_classes参数
self.num_classes = num_classes + 1  # 包含背景类

4. 训练策略与参数调优

DETR的训练需要特别注意学习率和批次大小的设置。对于小规模数据集，建议采用以下配置：

参数	推荐值	说明
lr	1e-4	基础学习率
lr_backbone	1e-5	骨干网络学习率
batch_size	2-4	取决于GPU显存
epochs	50-100	小数据集可适当减少
weight_decay	1e-4	权重衰减系数

启动训练的命令示例：

python main.py \
    --dataset_file "coco" \
    --coco_path "custom_dataset" \
    --epochs 100 \
    --lr=1e-4 \
    --batch_size=2 \
    --output_dir="outputs" \
    --resume="detr-r50-custom.pth"

常见训练问题及解决方案：

• loss不下降：尝试减小学习率或增加批次大小
• 显存不足：减小批次大小或输入图像尺寸
• 过拟合：增加数据增强或提前停止训练

5. 模型评估与结果可视化

训练完成后，可以使用DETR提供的工具进行结果可视化：

from util.plot_utils import plot_results
import matplotlib.pyplot as plt

# 加载模型和测试图像
model = torch.load("outputs/checkpoint.pth")
img = Image.open("test.jpg")

# 运行推理并可视化
outputs = model(img)
plot_results(img, outputs, threshold=0.7)
plt.show()

评估指标解读：

• AP (Average Precision): 综合考量精确率和召回率
• AP50: IoU阈值为0.5时的AP
• AP75: IoU阈值为0.75时的AP
• mAP: 多类别AP的平均值

6. 实际应用与性能优化

将训练好的模型部署到实际应用中时，可以考虑以下优化手段：

1. 模型量化：使用PyTorch的量化工具减小模型大小

   model_quantized = torch.quantization.quantize_dynamic(
       model, {torch.nn.Linear}, dtype=torch.qint8)
   

2. ONNX导出：转换为通用格式便于跨平台部署

   torch.onnx.export(model, dummy_input, "detr.onnx", 
       input_names=["input"], output_names=["output"])
   

3. TensorRT加速：针对NVIDIA GPU的优化推理

对于实时性要求高的场景，可以调整Decoder的层数或输入分辨率来平衡速度和精度。