从零复现DETR：环境搭建、核心模块与训练调参全记录

1. 为什么DETR值得你花时间亲手复现一遍如果你最近在目标检测这个方向上有过一些摸索大概率会频繁听到一个名字DETR。它的全称是DEtection TRansformer核心思路是把目标检测这个 traditionally 依赖锚框和非极大值抑制的任务直接转化成一个端到端的集合预测问题。听起来很抽象但实际跑一遍代码你就会发现它的设计哲学比想象中要干净利落得多。我第一次接触DETR是在一个需要做密集小目标检测的项目里。当时试过好几种基于锚框的方案调参调到怀疑人生——锚框尺寸、长宽比、正负样本匹配策略、NMS阈值每一个都像是一个需要反复试探的黑盒。后来看到DETR的论文第一反应是“这也太简洁了”第二反应是“但真的能work吗”。带着这个疑问我决定把官方代码拉下来从头到尾复现一遍。这篇文章就是那次复现过程的完整记录。我会从环境搭建开始一步步拆解DETR的核心模块包括骨干网络、Transformer编解码器、二分图匹配损失、以及训练和推理的完整流程。过程中会穿插我踩过的坑、调参的心得、以及一些官方文档里不会写的细节。无论你是刚入门目标检测的新手还是想从锚框方案迁移到端到端方案的老手这篇内容都能给你一个可以直接抄作业的参考。需要提前说明的是DETR的复现涉及PyTorch、torchvision、以及一些辅助库的配合对显存有一定要求。我建议至少准备一张8GB显存以上的显卡如果显存不够可以通过减小batch size或者输入分辨率来适配。另外复现过程中遇到版本不兼容是常态我会在对应章节给出我实测通过的版本组合。2. 环境搭建与依赖版本的血泪教训2.1 版本组合的选择逻辑DETR官方仓库的代码是基于PyTorch 1.5左右写的但如果你直接照搬那个版本会发现很多新的CUDA驱动和Python版本已经不支持了。我试过三种组合最终稳定跑通的是PyTorch 1.12.1 torchvision 0.13.1 CUDA 11.3这个搭配。为什么选这个组合因为PyTorch 1.12对Transformer相关的算子做了不少优化而且和torchvision的兼容性经过大量项目验证踩坑概率最低。如果你用的是更新的PyTorch 2.x版本理论上也能跑但需要注意几个点一是torch.nn.Transformer的默认行为在2.x里有细微变化二是混合精度训练相关的API有调整。我建议第一次复现时先用1.12这个稳定版本等跑通了再尝试升级。安装命令我直接列出来你可以按需调整conda create -n detr python3.9 conda activate detr pip install torch1.12.1cu113 torchvision0.13.1cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install pycocotools scipy matplotlib opencv-python这里有个细节pycocotools在Windows上安装经常报错如果你用的是Windows建议直接上WSL2或者Linux环境。我在Windows原生环境下折腾了两个小时没搞定换到Ubuntu之后五分钟就装好了。2.2 数据集准备中的隐藏坑点DETR官方用的是COCO数据集但完整版COCO有十几万张图训练一轮就要很久。我建议第一次复现时先用一个子集比如val2017里的前500张图配合对应的标注文件快速验证代码能不能跑通。等流程走通了再换全量数据。数据目录结构必须严格按照COCO的格式来组织data/ coco/ annotations/ instances_train2017.json instances_val2017.json train2017/ val2017/这里有个容易忽略的点DETR的dataset类在读取标注时会默认过滤掉iscrowd1的样本但如果你自己制作的数据集没有这个字段代码会直接报错。解决办法是在生成标注文件时给每个标注都加上iscrowd: 0。这个坑我踩过当时排查了半天才发现是字段缺失导致的。另外DETR的输入图像默认会缩放到短边800、长边不超过1333的尺寸。如果你的数据集里有很多极端长宽比的图建议提前做一下筛选或者裁剪否则训练时会出现显存爆炸的情况。2.3 预训练权重加载的注意事项官方提供了几个预训练模型包括ResNet-50和ResNet-101两种骨干网络。我建议从ResNet-50版本开始因为它的参数量相对小一些训练和推理速度都更快。下载权重之后加载时要注意key的匹配问题——如果你修改了模型结构比如改了类别数需要把分类头的权重过滤掉再加载。checkpoint torch.load(detr-r50.pth, map_locationcpu) model_dict model.state_dict() checkpoint {k: v for k, v in checkpoint.items() if k in model_dict and class_embed not in k} model_dict.update(checkpoint) model.load_state_dict(model_dict)这段代码的意思是只加载骨干网络和Transformer部分的权重分类头因为类别数不同需要重新初始化。如果你直接全量加载会报size mismatch的错误。3. DETR核心模块的拆解与代码实现3.1 骨干网络与位置编码的配合方式DETR的骨干网络用的是ResNet但和普通分类任务不同的是它只取到layer4的输出然后经过一个1x1卷积把通道数降到256。这个256就是整个Transformer的隐藏维度。为什么是256因为这是论文里经过实验验证的平衡点——再小的话表达能力不够再大的话计算量会急剧增加。位置编码是DETR里比较有意思的部分。它用的是正弦余弦编码但和原始Transformer不同的是DETR把位置编码加在了骨干网络输出的特征图上而不是加在输入序列上。具体来说对于一张H x W的特征图会生成一个同样大小的位置编码矩阵然后和特征图逐元素相加。class PositionEmbeddingSine(nn.Module): def __init__(self, num_pos_feats128, temperature10000): super().__init__() self.num_pos_feats num_pos_feats self.temperature temperature self.scale 2 * math.pi def forward(self, tensor): b, c, h, w tensor.shape y_embed torch.arange(h, dtypetorch.float32, devicetensor.device).unsqueeze(1).repeat(1, w) x_embed torch.arange(w, dtypetorch.float32, devicetensor.device).unsqueeze(0).repeat(h, 1) y_embed y_embed / (h 1e-6) * self.scale x_embed x_embed / (w 1e-6) * self.scale dim_t torch.arange(self.num_pos_feats, dtypetorch.float32, devicetensor.device) dim_t self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) pos_x x_embed[:, :, None] / dim_t pos_y y_embed[:, :, None] / dim_t pos_x torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim3).flatten(2) pos_y torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim3).flatten(2) pos torch.cat((pos_y, pos_x), dim2).permute(0, 3, 1, 2) return pos这段代码的关键在于scale的设定和dim_t的计算。scale设为2*pi是为了让位置编码覆盖一个完整的周期dim_t则决定了不同维度的频率。实测下来这个配置在COCO上表现很稳定不需要额外调整。3.2 Transformer编解码器的结构细节DETR的Transformer部分包括一个编码器和一个解码器。编码器由6层标准Transformer层组成每层包括自注意力、前馈网络、残差连接和层归一化。解码器也是6层但多了一个交叉注意力模块用来让查询向量关注编码器的输出。这里有个关键设计解码器的输入是一组可学习的查询向量数量固定为100。这100个查询向量经过解码器之后每个都会输出一个预测结果包括类别和边界框坐标。这也是DETR不需要NMS的原因——每个查询向量负责预测一个目标它们之间通过二分图匹配损失来保证不重复。class Transformer(nn.Module): def __init__(self, d_model256, nhead8, num_encoder_layers6, num_decoder_layers6, dim_feedforward2048, dropout0.1): super().__init__() self.encoder TransformerEncoder( TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout), num_encoder_layers ) self.decoder TransformerDecoder( TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout), num_decoder_layers ) self.query_embed nn.Embedding(100, d_model) def forward(self, src, pos_embed): bs, c, h, w src.shape src src.flatten(2).permute(2, 0, 1) pos_embed pos_embed.flatten(2).permute(2, 0, 1) query_embed self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1) tgt torch.zeros_like(query_embed) memory self.encoder(src, pospos_embed) hs self.decoder(tgt, memory, pospos_embed, query_posquery_embed) return hs.transpose(1, 2)dim_feedforward设为2048是d_model的8倍这个比例在Transformer类模型里比较常见。dropout设为0.1是为了防止过拟合。实测下来如果数据集比较小可以把dropout调到0.2甚至0.3效果会更稳。3.3 二分图匹配损失的实现逻辑二分图匹配损失是DETR最核心的创新点之一。简单来说对于每一张图模型会输出100个预测结果而真实标注可能只有几个目标。我们需要找到一个最优的匹配方案让每个真实目标都对应到一个唯一的预测结果同时让匹配的总代价最小。这个匹配过程用的是匈牙利算法代价函数包括分类代价和边界框代价两部分。分类代价用的是交叉熵边界框代价用的是L1损失和GIoU损失的加权和。class HungarianMatcher(nn.Module): def __init__(self, cost_class1, cost_bbox5, cost_giou2): super().__init__() self.cost_class cost_class self.cost_bbox cost_bbox self.cost_giou cost_giou torch.no_grad() def forward(self, outputs, targets): bs, num_queries outputs[pred_logits].shape[:2] out_prob outputs[pred_logits].flatten(0, 1).softmax(-1) out_bbox outputs[pred_boxes].flatten(0, 1) tgt_ids torch.cat([v[labels] for v in targets]) tgt_bbox torch.cat([v[boxes] for v in targets]) cost_class -out_prob[:, tgt_ids] cost_bbox torch.cdist(out_bbox, tgt_bbox, p1) cost_giou -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) C self.cost_bbox * cost_bbox self.cost_class * cost_class self.cost_giou * cost_giou C C.view(bs, num_queries, -1).cpu() sizes [len(v[boxes]) for v in targets] indices [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))] return [(torch.as_tensor(i, dtypetorch.int64), torch.as_tensor(j, dtypetorch.int64)) for i, j in indices]cost_class、cost_bbox、cost_giou这三个权重分别是1、5、2这是论文里给出的推荐值。我试过调整这些权重发现对最终精度的影响其实不大但如果把cost_giou调得太低边界框的回归效果会明显变差。所以建议第一次复现时直接用默认值等跑通了再尝试微调。4. 训练流程的完整跑通与调参心得4.1 训练脚本的关键参数设置DETR的训练脚本里有几个参数需要特别注意。首先是学习率骨干网络和Transformer部分用的学习率是不一样的——骨干网络用1e-5Transformer用1e-4。这是因为骨干网络是预训练好的不需要太大的学习率来调整而Transformer是从头训练的需要更大的学习率来快速收敛。param_dicts [ {params: [p for n, p in model.named_parameters() if backbone not in n and p.requires_grad]}, {params: [p for n, p in model.named_parameters() if backbone in n and p.requires_grad], lr: 1e-5}, ] optimizer torch.optim.AdamW(param_dicts, lr1e-4, weight_decay1e-4)权重衰减设为1e-4用的是AdamW而不是Adam因为AdamW对权重衰减的处理更合理。另外学习率调度用的是StepLR在第100个epoch时衰减10倍。如果你训练的数据集比较小可以把衰减点提前到第50个epoch。Batch size方面官方用的是16但这需要至少16GB显存。如果显存不够可以降到8或者4同时把学习率按比例调小。我试过batch size4的情况最终精度会掉1-2个点但训练时间会缩短很多适合快速验证。4.2 损失曲线的解读与异常判断训练过程中损失曲线是最直观的监控指标。DETR的总损失包括分类损失、L1边界框损失、GIoU损失三部分。正常情况下这三部分损失都会随着训练轮次增加而稳步下降。如果发现某一项损失不降反升通常意味着对应的权重设置有问题。我遇到过一种情况分类损失一直在下降但GIoU损失卡在某个值不动了。排查后发现是学习率太大导致边界框回归的参数在最优值附近震荡。把Transformer的学习率从1e-4降到5e-5之后GIoU损失就正常下降了。另一个常见的异常是损失突然变成NaN。这通常是因为梯度爆炸解决办法是加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm0.1)max_norm设为0.1这个值在DETR的训练里比较常用。如果还是出现NaN可以进一步降到0.01。4.3 学习率调度的实操经验DETR原论文用的是StepLR在第100个epoch衰减一次。但我在实际训练中发现如果数据集比较小比如只有几千张图模型往往在第50个epoch左右就已经收敛了继续训练只会导致过拟合。这时候可以把衰减点提前或者改用CosineAnnealingLR让学习率平滑下降。scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100)CosineAnnealingLR的好处是不需要手动设置衰减点学习率会按照余弦曲线从初始值降到0。实测下来这种调度方式在小数据集上的表现比StepLR更稳定最终精度也能高0.5个点左右。另外DETR对学习率预热比较敏感。前几个epoch如果学习率太大模型很容易发散。建议加一个线性预热def warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor): def f(x): if x warmup_iters: return 1 alpha float(x) / warmup_iters return warmup_factor * (1 - alpha) alpha return torch.optim.lr_scheduler.LambdaLR(optimizer, f)warmup_iters设为1000左右warmup_factor设为0.1。这样前1000次迭代的学习率会从0.1倍逐渐升到正常值有效避免早期发散。5. 推理与评估环节的实操细节5.1 后处理流程的简化与注意事项DETR的推理输出是100个预测结果每个结果包括类别概率和边界框坐标。后处理的第一步是过滤掉低置信度的预测通常阈值设为0.7左右。然后对剩下的预测做类别判断取概率最大的类别作为最终类别。def postprocess(outputs, threshold0.7): prob outputs[pred_logits].softmax(-1)[..., :-1] scores, labels prob.max(-1) boxes outputs[pred_boxes] keep scores threshold scores scores[keep] labels labels[keep] boxes boxes[keep] return scores, labels, boxes这里有个细节pred_logits的最后一维是类别数加1多出来的那个是“无目标”类别。所以在取max之前要把最后一列去掉。这个操作在官方代码里是通过[..., :-1]实现的但如果你自己改过类别数记得同步调整。另外DETR输出的边界框是归一化的中心点坐标和宽高格式是cxcywh。如果需要转换成xyxy格式可以用下面的函数def box_cxcywh_to_xyxy(x): x_c, y_c, w, h x.unbind(-1) b [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c 0.5 * w), (y_c 0.5 * h)] return torch.stack(b, dim-1)这个转换在计算IoU或者可视化的时候会用到建议提前封装好。5.2 COCO评估指标的计算与解读DETR的评估用的是COCO的官方指标包括AP、AP50、AP75、APs、APm、APl等。其中AP是IoU阈值从0.5到0.95的平均值AP50是IoU0.5时的精度AP75是IoU0.75时的精度。APs、APm、APl分别对应小、中、大目标的精度。from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval coco_gt COCO(annotations/instances_val2017.json) coco_dt coco_gt.loadRes(predictions.json) coco_eval COCOeval(coco_gt, coco_dt, bbox) coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize()summarize()会输出12个指标其中第一个就是AP。官方在COCO上报告的AP是42.0左右如果你复现出来的结果在40以上说明流程基本正确。如果低于35可能是某个环节出了问题建议从数据加载、损失计算、后处理这三个地方逐一排查。我复现时第一次跑出来的AP只有32排查后发现是学习率预热没加导致模型早期发散。加上预热之后AP直接升到了41.5。所以这些细节对最终结果的影响真的很大不能偷懒跳过。5.3 可视化验证的实用技巧光看数字指标有时候不够直观建议把预测结果画到图上肉眼检查一下。DETR的预测框通常比较准确但偶尔会出现重复预测的情况——虽然理论上二分图匹配应该避免这个问题但在实际推理时如果两个查询向量都预测了同一个目标还是会出现重复框。import matplotlib.pyplot as plt import matplotlib.patches as patches def visualize(image, boxes, labels, scores): fig, ax plt.subplots(1) ax.imshow(image) for box, label, score in zip(boxes, labels, scores): x1, y1, x2, y2 box rect patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth2, edgecolorr, facecolornone) ax.add_patch(rect) ax.text(x1, y1, f{label}: {score:.2f}, colorwhite, fontsize12, bboxdict(facecolorred, alpha0.5)) plt.show()可视化的时候注意把图像从归一化坐标还原到原始像素坐标否则框的位置会不对。还原公式是x * width和y * height其中width和height是原始图像的尺寸。6. 复现过程中最容易卡住的几个问题6.1 显存不足的多种应对策略显存不足是复现DETR时最常见的问题。官方配置需要16GB显存但很多人手头只有8GB甚至6GB的卡。除了减小batch size还有几个办法可以试试。第一个是减小输入分辨率。DETR默认把图像短边缩放到800你可以改成600甚至400。这样特征图的尺寸会变小显存占用会显著降低。代价是小目标的检测精度会下降但如果你的数据集里没有太多小目标影响不大。第二个是使用梯度累积。把batch size设为4然后每4次迭代才更新一次参数等效于batch size16。这样显存占用只有原来的四分之一但训练时间会变长。accumulation_steps 4 for i, (images, targets) in enumerate(dataloader): outputs model(images) loss criterion(outputs, targets) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()第三个是使用混合精度训练。PyTorch的amp模块可以自动把部分计算转成float16显存占用能降低30%左右。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(images) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这三个方法可以组合使用我试过在6GB显存的卡上跑通DETR靠的就是batch size2 分辨率600 混合精度。6.2 训练不收敛的排查链路如果训练了几十个epoch损失还是不降可以按照下面的顺序逐一排查。第一步检查数据加载是否正确。把dataloader里的一个batch取出来可视化一下图像和标注确认标注框的位置和类别都是对的。我遇到过标注文件里坐标是xyxy格式但代码按xywh解析的情况导致框的位置完全错乱。第二步检查损失函数的输入输出。把模型输出和真实标注打印出来看看形状和数值范围是否合理。特别是分类损失如果真实标注的类别索引超出了模型输出的类别数会直接报错或者算出错误的结果。第三步检查学习率是否合适。如果学习率太大损失会震荡甚至发散如果太小损失下降会非常慢。可以先用一个很小的学习率比如1e-5跑几个epoch确认损失在下降然后再逐步调大。第四步检查预训练权重是否加载成功。如果骨干网络的权重没加载上模型相当于从头训练收敛会慢很多。可以在加载权重后打印几个参数的值和原始权重对比一下。6.3 评估指标偏低的常见原因如果训练损失正常下降但评估指标偏低通常有以下几个原因。一是后处理阈值设得不对。默认的0.7可能不适合你的数据集可以试试0.5或者0.8看哪个指标更高。二是类别不平衡。如果某些类别的样本特别少模型可能会偏向预测那些样本多的类别。解决办法是对损失函数里的分类损失加权重或者对少样本类别做数据增强。三是边界框回归不够准。可以检查一下GIoU损失的权重是不是设得太低适当调高可能会改善边界框的精度。四是过拟合。如果训练损失很低但验证损失很高说明模型过拟合了。可以增加dropout、加权重衰减、或者用更多的数据增强。7. 从复现到改进的几条可行路径7.1 针对小目标检测的改进思路DETR在小目标上的表现一直是个短板主要原因是骨干网络输出的特征图分辨率比较低通常是原图的1/32小目标在特征图上可能只有几个像素。一个常见的改进思路是引入多尺度特征比如用FPN结构把不同尺度的特征融合起来。具体做法是在骨干网络后面接一个FPN输出四个尺度的特征图然后分别做位置编码再拼接成一个更长的序列输入Transformer。这样小目标可以在高分辨率特征图上被更好地捕捉到。我试过这个改法在小目标上的AP能提升3-5个点但计算量会增加不少。另一个思路是提高输入分辨率。把短边从800提到1200甚至1600小目标的像素数会显著增加。代价是显存占用和推理时间都会上升需要根据实际场景权衡。7.2 训练效率优化的几个方向DETR的训练速度比较慢官方配置下训练300个epoch需要好几天。如果只是想快速验证想法可以试试以下几个优化方向。一是用更小的骨干网络。把ResNet-50换成ResNet-18参数量减少一半以上训练速度能提升不少。代价是精度会掉几个点但作为快速验证是够用的。二是减少解码器层数。把6层降到3层计算量会明显降低。实测下来精度掉得不多但训练速度能快30%左右。三是用更高效的注意力实现。PyTorch 2.0之后引入了scaled_dot_product_attention比手动实现的注意力快不少。如果你的PyTorch版本支持可以把Transformer里的注意力替换成这个。7.3 迁移到自定义数据集的注意事项如果你想把DETR迁移到自己的数据集上有几个地方需要调整。首先是类别数。把num_classes改成你的数据集类别数加1加1是因为有无目标类别。同时分类头的输出维度也要同步修改。其次是标注格式。DETR期望的标注是cxcywh格式的归一化坐标如果你的标注是xyxy格式的绝对坐标需要先转换。最后是数据增强。DETR原论文用了随机裁剪、随机缩放、颜色抖动等增强手段。如果你的数据集比较小建议把这些增强都加上能有效防止过拟合。transform T.Compose([ T.RandomHorizontalFlip(), T.RandomResize([600, 800], max_size1000), T.ColorJitter(0.4, 0.4, 0.4, 0.1), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])RandomResize的短边范围设为600到800长边不超过1000。这个配置在大多数数据集上都能work不需要太多调整。我在实际迁移到自定义数据集时遇到过一个比较隐蔽的问题数据集中有些图像的尺寸特别小比如只有几十个像素。这些图像经过缩放后特征图上的目标几乎消失了导致模型完全学不到。解决办法是在数据加载时过滤掉这些极端样本或者对它们做上采样。整体走下来DETR的复现难度属于中等偏上主要难点在于环境配置和训练调参。但一旦跑通你会发现它的代码结构非常清晰扩展性也很好。后续如果想尝试Deformable DETR、DAB-DETR这些改进版本有了这次复现的基础上手会快很多。