DETR目标检测实战：从零开始训练自己的Transformer模型（附完整代码）

DETR目标检测实战从零开始训练自己的Transformer模型附完整代码在计算机视觉领域目标检测一直是最具挑战性的任务之一。传统基于卷积神经网络(CNN)的方法虽然取得了显著成果但往往需要复杂的后处理流程和手工设计的组件。2020年Facebook AI团队提出的DETR(DEtection TRansformer)彻底改变了这一局面首次将Transformer架构成功应用于目标检测任务实现了端到端的检测流程。本文将带您从零开始完整实现一个DETR模型的训练过程。1. 环境准备与依赖安装训练DETR模型前我们需要搭建合适的开发环境。推荐使用Python 3.8和PyTorch 1.7版本这是经过验证的稳定组合。核心依赖包清单pip install torch1.8.0cu111 torchvision0.9.0cu111 -f https://download.pytorch.org/whl/torch_stable.html pip install pycocotools matplotlib scipy opencv-python对于GPU加速确保正确配置了CUDA环境。可以通过以下命令验证import torch print(torch.cuda.is_available()) # 应返回True print(torch.cuda.device_count()) # 显示可用GPU数量常见问题排查如果遇到CUDA版本不兼容建议使用conda管理环境Windows用户可能需要额外安装Visual C构建工具对于内存有限的GPU可尝试降低batch size或使用梯度累积2. 数据集准备与预处理DETR支持标准的目标检测数据集格式最常用的是COCO格式。我们以COCO 2017数据集为例展示如何准备训练数据。2.1 数据集结构标准的COCO数据集目录结构如下coco/ ├── annotations │ ├── instances_train2017.json │ └── instances_val2017.json ├── train2017 │ └── *.jpg └── val2017 └── *.jpg2.2 数据增强策略DETR对数据增强相对敏感推荐使用以下组合from torchvision import transforms train_transforms transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.RandomResizedCrop(size(800, 800), scale(0.8, 1.0)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) val_transforms transforms.Compose([ transforms.Resize(800), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])注意DETR对输入图像尺寸没有严格要求但建议保持长宽比接近1:1800x800是一个经过验证的合理尺寸。3. 模型构建与配置DETR的核心架构包含四个主要组件CNN骨干网络、Transformer编码器、Transformer解码器和预测头。3.1 骨干网络实现DETR默认使用ResNet-50作为骨干网络import torchvision from torch import nn class Backbone(nn.Module): def __init__(self, nameresnet50, train_backboneTrue, dilationFalse): super().__init__() backbone getattr(torchvision.models, name)( replace_stride_with_dilation[False, False, dilation], pretrainedTrue) self.num_channels 2048 if name in (resnet50, resnet101) else 512 self.body nn.Sequential( backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool, backbone.layer1, backbone.layer2, backbone.layer3, backbone.layer4 ) def forward(self, x): return self.body(x)3.2 Transformer实现DETR使用标准的Transformer架构但有一些关键修改from torch.nn import MultiheadAttention class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward2048, dropout0.1): super().__init__() self.self_attn MultiheadAttention(d_model, nhead, dropoutdropout) self.linear1 nn.Linear(d_model, dim_feedforward) self.dropout nn.Dropout(dropout) self.linear2 nn.Linear(dim_feedforward, d_model) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) self.activation nn.ReLU()3.3 完整模型组装将各组件整合为完整DETR模型class DETR(nn.Module): def __init__(self, num_classes, num_queries100, hidden_dim256): super().__init__() self.backbone Backbone() self.conv nn.Conv2d(2048, hidden_dim, 1) self.transformer Transformer(d_modelhidden_dim) self.query_embed nn.Embedding(num_queries, hidden_dim) self.class_embed nn.Linear(hidden_dim, num_classes 1) self.bbox_embed MLP(hidden_dim, hidden_dim, 4, 3) def forward(self, x): features self.backbone(x) features self.conv(features) hs self.transformer(features, self.query_embed.weight) outputs_class self.class_embed(hs) outputs_coord self.bbox_embed(hs).sigmoid() return {pred_logits: outputs_class[-1], pred_boxes: outputs_coord[-1]}4. 训练流程与技巧DETR的训练有其特殊性需要特别注意损失函数和优化策略。4.1 匈牙利匹配损失实现DETR的核心创新之一是基于匈牙利算法的二部图匹配损失from scipy.optimize import linear_sum_assignment def hungarian_matcher(outputs, targets): bs, num_queries outputs[pred_logits].shape[:2] indices [] for i in range(bs): cost_class -out_prob[i] cost_bbox torch.cdist(out_bbox[i], tgt_bbox[i], p1) cost_giou -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox[i]), box_cxcywh_to_xyxy(tgt_bbox[i])) C cost_bbox cost_class cost_giou C C.reshape(num_queries, -1).cpu() indices.append(linear_sum_assignment(C)) return indices4.2 训练超参数设置经过实验验证的推荐参数参数推荐值说明学习率1e-4使用线性warmupBatch Size4-8取决于GPU内存Epochs300需要长时间训练优化器AdamW权重衰减0.0001学习率调度StepLR每200epoch衰减10倍4.3 训练加速技巧混合精度训练使用Apex或PyTorch原生AMPfrom torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()梯度累积在小批量GPU上模拟大批量训练accumulation_steps 4 for i, (inputs, targets) in enumerate(dataloader): loss criterion(model(inputs), targets) / accumulation_steps loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()5. 模型评估与可视化训练完成后我们需要评估模型性能并进行结果分析。5.1 评估指标实现标准COCO评估指标实现from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval def evaluate(model, dataset, device): model.eval() results [] for img, _ in dataset: with torch.no_grad(): outputs model(img.unsqueeze(0).to(device)) # 转换输出为COCO格式 # ... 省略转换代码 ... results.extend(coco_results) coco_gt COCO(dataset.coco) coco_dt coco_gt.loadRes(results) coco_eval COCOeval(coco_gt, coco_dt, bbox) coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() return coco_eval.stats5.2 结果可视化使用Matplotlib绘制检测结果import matplotlib.pyplot as plt def plot_results(img, outputs, threshold0.7): probas outputs[pred_logits].softmax(-1)[0, :, :-1] keep probas.max(-1).values threshold fig, ax plt.subplots(1, figsize(10,10)) ax.imshow(img) for p, (x, y, w, h) in zip(probas[keep], outputs[pred_boxes][0, keep]): ax.add_patch(plt.Rectangle((x-w/2, y-h/2), w, h, fillFalse, colorr, linewidth2)) cl p.argmax() text f{CLASSES[cl]}: {p[cl]:0.2f} ax.text(x-w/2, y-h/2, text, fontsize10, bboxdict(facecoloryellow, alpha0.5)) plt.axis(off) plt.show()在实际项目中DETR的训练往往需要3-5天时间在8块V100上但得到的模型具有很好的泛化能力。一个常见的问题是初期训练损失下降缓慢这属于正常现象通常50个epoch后会看到明显提升。