1. 为什么从YOLO转向DETR:端到端目标检测的本质变化
干了几年目标检测,YOLO系列我确实用过不少,从最初的基础版本一路看到后续各种改进版本。说实话,YOLO的工作流大家太熟了:锚框设计、后处理NMS、正负样本分配,每一步都有讲究,调参空间大,细节多。但这类模型的问题也很明显——后处理逻辑复杂,训练时各种技巧堆叠,部署时又得合并分支。
直到我实际动手跑通了DETR(Detection Transformer),才意识到检测框架还能这么设计。它把目标检测直接定义成一个集合预测问题,输入图像,输出一组预测框和对应类别,没有锚框、没有NMS、没有Region Proposal,整个pipeline干净利落。你训练完了直接推理就行,不需要额外的后处理模块,这对工程落地来说省了太多事。
DETR的核心思路其实很简单:把一张图喂给CNN骨干网络提取特征,展平后加上位置编码,送到Transformer的Encoder-Decoder结构里,最后通过一个检测头输出固定数量的预测框。这些框和真实的物体框做二分图匹配,直接计算损失,一步到位。
当初第一次听到这个思路,我也怀疑过:Transformer这种东西用在目标检测上,真的能稳吗?实测下来,它在COCO数据集上的AP能达到和当时Faster R-CNN相当的水平,而且在大物体检测上表现出明显的优势。虽然小物体检测还是弱项,但它的简洁性和模块化设计,确实给检测领域带来了一个新的方向。
这篇文章要分享的就是DETR从零到一的全流程:怎么搭环境、怎么准备数据、怎么训练和评估、怎么用训好的模型做预测。不管你是从YOLO阵营转过来的,还是第一次接触Transformer系检测模型,照着这套流程,你至少能少踩一半的坑。
1.1 标题拆解:为什么强调“告别YOLO”而不是“超越YOLO”
这里先澄清一点。写“告别YOLO”不是说DETR已经全面碾压YOLO了,而是说检测范式上,DETR代表了一种完全不同的解决思路。YOLO系是anchor-based或anchor-free的密集预测,空间网格上做分类和回归,再把重复框压掉;DETR是sparse的,一次只输出固定数量(比如100个)的查询,每个查询通过Transformer去“注意”图中某个物体区域,输出一个框。
这两种思路落实到套路层面差别很大:
- 训练流程:YOLO训练要看很多trick,比如预热、多尺度、Mosaic增强;DETR是端到端训练,直接用集合匹配损失,训练稳定性和收敛节奏相对可控。
- 推理流程:YOLO推理后必须加NMS去重;DETR直接输出最终结果,整个后处理几乎等于零。
- 工程链路:DETR的pipeline更短,做部署和二次开发时逻辑更清晰。
所以“告别YOLO”的准确含义是:如果你手头的场景允许你换框架,DETR这条路线值得认真跑一遍,它解决了很多传统检测框架中人工设计的复杂度。尤其当你做研究、做原型验证,或者需要频繁迭代检测模型结构时,DETR这种化繁为简的做法收益很大。
1.2 适合谁看以及需要的前置知识
这篇内容的阅读对象大概有三类:
第一类是已经熟悉YOLO系训练流程,但想拓展视野的工程师。你不需要重新学一堆锚框、NMS之类的概念,DETR这套路线几乎是反着来的。
第二类是刚入门目标检测的学生或研究者,想找一个基线模型做实验。DETR虽然是2020年的工作,但它作为Transformer检测器的起点,结构清晰,代码量适中,很适合做二次开发的基座。
第三类是想把Transformer用到视觉任务上、但不知道从哪里下手的开发者。
前置知识方面,最低要求是:跑过至少一个目标检测训练脚本,懂一点PyTorch基本功,知道什么是backbone、什么是loss。Transformer的原理你可以不深究——我在下面会尽量讲人话——但你至少得知道self-attention大概是做特征交互的东西。如果完全不懂Transformer也没关系,先装上跑通,再回头理解内部结构,也行。
2. 环境搭建:版本搭配与依赖安装细节
DETR的官方代码基于PyTorch,整体依赖不算多,但版本敏感。我一开始图省事直接装最新版PyTorch,结果编译torchvision的扩展模块时报了一堆错。折腾半天才发现是CUDA版本和编译工具链不匹配的问题。
先说结论,我最终稳定跑通的组合是:
- Python:3.8或3.10都行,但3.8更稳一些。
- PyTorch:1.7以上即可,推荐1.10或2.x对应版本。
- CUDA:10.2或11.x,看驱动和显卡。官方代码本身对CUDA版本要求不高,关键是torch和torchvision的对应关系要对。
- torchvision:0.8.0及以上,因为DETR里要用到torchvision的算子。
- 其他依赖:scipy,另外跑可视化或demo时用opencv-python。
如果你用的是新一些的PyTorch 2.x版本,也没问题,很多社区版DETR都已经适配了。但是要注意,官方仓库里一些编译扩展,比如MultiScaleDeformableAttention(那是Deformable DETR的内容,原版DETR没有),如果你跑的是原版DETR,就不用管这些编译的事。
我建议新建一个独立的conda环境,别跟其他项目混在一起。因为torch、torchvision、cuda版本一混,排查问题的时间远超你重装环境的时间。
2.1 安装步骤:三步走,缺一不可
第一步,创建环境并安装PyTorch。
conda create -n detr python=3.8 conda activate detr # 根据你自己的CUDA版本选择命令,下面是以CUDA 11.3为例 pip install torch==1.12.1 torchvision==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu113第二步,克隆官方仓库并安装依赖。
git clone https://github.com/facebookresearch/detr.git cd detr pip install scipy pip install -r requirements.txtrequirements.txt里就是一些基础库,没有奇怪的依赖,一般不会出问题。
第三步,验证安装。
python -c "import torch; print(torch.__version__, torch.cuda.is_available())" python -c "import torchvision; print(torchvision.__version__)"这一步必须看到cuda.is_available()返回True,否则后面训练会慢到怀疑人生,或者直接报内存错误。
有些读者可能卡在torchvision的编译上。如果你用的是pip直接安装的torchvision,基本不需要编译;如果你从源码安装,那就需要确认gcc版本和CUDA toolkit匹配。我的建议是:能用pip装就绝不要源码编译,除非你确实需要改到torchvision里面的源码。
2.2 依赖版本对照速查表
为了减少大家试错的时间,我把常见可用搭配整理成了下面的表格:
| Python版本 | torch版本 | torchvision版本 | CUDA版本 | 实测状态 |
|---|---|---|---|---|
| 3.8 | 1.7.0 | 0.8.1 | 10.2 | 稳定运行 |
| 3.8 | 1.12.1 | 0.13.1 | 11.3 | 稳定运行 |
| 3.10 | 2.0.0 | 0.15.1 | 11.7 | 稳定运行,需微调代码 |
| 3.10 | 2.1.0 | 0.16.0 | 11.8 | 稳定运行 |
最后一个组合需要注意:PyTorch 2.x在DETR原版代码上可能会出现一些API废弃警告,但不影响训练。如果你看到类似torch.nn.functional.sigmoid相关的弃用提示,忽略即可。
安装完之后,我建议顺手跑一个官方提供的预训练模型做一次推理,确认整体链路没问题,再开始训练自己的数据。这么做的好处是:能区分“环境问题”和“代码问题”,后面训练出了问题不会两头找。
3. 准备数据:COCO格式与训练前的关键检查
DETR官方代码默认使用COCO数据集格式,虽然你可以自己写Dataset类来适配其他格式,但对第一次跑的人来说,建议直接用COCO格式,能省掉大量调试时间。
你可能已经熟悉了YOLO那种txt标签格式,每一行是类别、中心x、中心y、宽、高。DETR这套代码走的是COCO的JSON结构。区别在于:COCO的标注是绝对像素坐标,也就是框的左上角x、左上角y、宽、高,而且类别ID是从1开始编号的,0留给背景。这个细节一定要注意,我第一次转换数据时就因为类别ID从0开始,导致训练出来的模型所有类别全部错位,检测结果完全没法看。
3.1 COCO格式一次性讲清楚
COCO数据集的核心是一个JSON文件,里面主要包括以下几个字段:
- images:图像列表,每个元素有id、file_name、width、height。
- annotations:标注列表,每个元素有image_id、category_id、bbox([x, y, width, height])、area、iscrowd。
- categories:类别列表,每个元素有id、name。
DETR代码读取JSON后,内部会构建一个将标注映射到图像索引的结构,然后为每张图生成一个目标字典,包括boxes(格式为[cx, cy, w, h],且归一化到0到1)、labels、image_id等。
你自己转换数据时,写个脚本就行。要注意的几个点:
- bbox坐标必须是像素级,且不能有小数误差太大。
- 所有图像尺寸必须大于等于啥?没有硬性要求,但如果图像太小,会被DETR内部的下采样影响,导致细节丢失。建议最短边不要小于32像素。
- iscrowd标记为1的标注会被忽略或者在训练时有特殊处理,如果你自己的数据里没有crowd这个概念,直接全给0。
我习惯的做法是,把已有的XML或txt标注先用脚本统一转成COCO格式的JSON,然后写一个验证脚本,读取JSON里每条bbox,检查是否越界、宽高是否为负、类别ID是否在合理范围内。这一步看着麻烦,但能帮你避免训练时莫名奇妙报NaN或者loss不下降的问题。
3.2 自建数据集的最小脚本示例
这里分享一个极简转换脚本的思路。假设你的标注存在一个字典里,键是图片名,值是一个列表,每个元素是[x1, y1, x2, y2, label_id],其中label_id已是COCO格式的类别ID:
import json import os from PIL import Image def convert_to_coco(img_dir, annotations_dict, output_path): images = [] annotations = [] categories = [{"id": 1, "name": "cat"}, {"id": 2, "name": "dog"}] # 举例 ann_id = 1 for img_id, (img_name, anns) in enumerate(annotations_dict.items()): img_path = os.path.join(img_dir, img_name) w, h = Image.open(img_path).size images.append({ "id": img_id, "file_name": img_name, "width": w, "height": h }) for x1, y1, x2, y2, label_id in anns: bw = x2 - x1 bh = y2 - y1 annotations.append({ "id": ann_id, "image_id": img_id, "category_id": label_id, "bbox": [x1, y1, bw, bh], "area": bw * bh, "iscrowd": 0 }) ann_id += 1 coco_format = { "images": images, "annotations": annotations, "categories": categories } with open(output_path, "w") as f: json.dump(coco_format, f) print(f"Converted {len(images)} images, {len(annotations)} annotations")转换完以后,你能在COCO官方评测脚本或者DETR自带的datasets模块里直接跑通数据加载,才算合格。
3.3 训练集与验证集的划分策略
官方训练脚本只需要传入两个路径:--coco_path指向包含train2017和val2017子目录的数据根目录。你自建数据时,可以另建目录结构:
data/ ├── train/ │ ├── images/ │ └── annotations.json └── val/ ├── images/ └── annotations.json不过原版代码的COCODataset实现是硬编码加载train2017这种目录名的,所以如果你想省事,就直接按COCO的目录结构来——建一个coco父目录,里面放train2017和val2017,还有对应的annotations/instances_train2017.json和instances_val2017.json。这是最简单、最少改代码的方式。
如果你会改代码,也可以继承torch.utils.data.Dataset自己写一个,但第一次跑通之前,我不建议动原版数据接口,因为DETR的数据加载里有一些细节,比如NestedTensor的处理,改不好会引入隐蔽bug。
4. 训练与评估:从命令行到Loss曲线的完整过程
数据准备好了,环境也验证过预训练推理了,接下来就是核心环节:训练自己的模型。
DETR官方提供了main.py,启动训练只需要一行命令,但参数细节决定了训练是顺利还是折磨。
4.1 训练一个自定义小模型的完整命令
我这里以一个自定义的、只有少量类别的数据集为例(比如2类或者5类),给你一条实用的训练命令:
python main.py \ --dataset_file coco \ --coco_path ./data/coco \ --output_dir ./outputs/my_detr \ --resume ./outputs/my_detr/checkpoint.pth \ --epochs 50 \ --lr_drop 40 \ --batch_size 2 \ --num_workers 4 \ --num_queries 100解释一下关键参数:
--dataset_file coco:指定数据集格式。即使你用的是自建COCO格式数据,这里也填coco。--coco_path:数据根目录。--output_dir:模型输出目录,日志和checkpoint都存这里。--resume:从某个checkpoint继续训练。第一次跑时没有就删掉这行。--epochs:训练轮数。自建小数据集30到50轮比较常见。--lr_drop:学习率衰减的轮数。比如50轮训练,在40轮时衰减一次。--batch_size:batch size。DETR比较吃显存,默认backbone是ResNet-50,8卡跑batch size 2,单卡可能得开梯度累积或者用小模型。--num_queries:Transformer解码器输出的查询数量,默认100。如果你的场景里单张图目标数量很少,可以减到50甚至30,能略微提速,但AP可能会有轻微下降。
这里有个容易忽略的点:官方模型默认把backbone冻结的前5个阶段(stem和layer1)设为不参与训练?不对,实际上官方代码里backbone默认从头到尾全部参与训练的。如果你想缩短训练时间,可以加一个--backbone resnet50,同时设置--dilation和--position_embedding等参数。更省事的做法是加载ImageNet的预训练backbone权重,而--pretrained参数加载的是完整DETR模型权重,如果你只想用backbone预训练,需要自己改加载逻辑。
第一次实验我建议:直接用官方提供的ResNet-50 backbone,预训练权重从torchvision的模型库自动下载。这样虽然训练时间还是长,但至少收敛很快。
4.2 训练中显存不够怎么办
DETR显存占用高是出了名的,尤其是Decoder端,100个查询跟特征图做cross-attention,计算量不小。
如果batch size调到1,显存还是不够,优先考虑:
- 减小输入分辨率。官方默认
--coco_path那套用的是shortest_edge从480到800的随机缩放。你可以改成固定size,比如把最长边限制在640,能明显减少显存和训练时间。 - 用
--backbone resnet18或者--backbone resnet34。官方代码不一定支持直接传resnet18,需要改一下模型构建部分,但改动不大。 - 使用梯度累积。原代码没直接支持,但你可以写一个简单的loss累积逻辑,或直接外部套一个GradientAccumulator。
- 混精度训练。官方代码没有直接开AMP,需要自己加autocast和GradScaler。如果你用PyTorch 2.x,这块代码写起来非常轻松。
我自己在单张显卡上训练一个小数据集时,用的配置是:backbone为ResNet-18,输入最长边640,batch size为2,开了梯度累积4步,等效batch size是8。最终训练出来的模型在验证集上大概有41的mAP,当然这里说的是自己数据的相对指标,跟COCO那个mAP没有可比性。
4.3 Loss曲线怎么看
训练日志里每一项log都有含义。DETR的loss由三部分组成:分类loss、L1框回归loss、GIoU loss。
分类loss走的是focal loss,负责让模型判断每个查询对应的物体类别。L1 loss直接监督框的中心点坐标和宽高偏差。GIoU loss负责让预测框和真实框的重合度尽量高。三者权重默认分别是1、5、2,这个比例是官方调好的,一般不要动。
看loss曲线时,重点关注总loss有没有稳步下降,以及分类loss和GIoU loss的趋势是否协调。如果分类loss降得很慢,但GIoU降得快,说明匹配可能有问题;反之也一样。
还有一个很常见的现象:训练初期总loss会有一个小高峰,然后才开始下降。这是因为DETR里的匈牙利匹配在训练开始阶段不稳定,查询还没有学会聚焦到具体的物体区域,导致匹配关系剧烈变化,loss自然波动大。如果你看到前几个epoch loss上下波动甚至轻微上升,不用慌,让它跑下去,一般在第5到第10个epoch后会开始稳定下降。
4.4 验证与评估:mAP怎么算
训练结束后,用main.py的--eval参数即可评估:
python main.py \ --dataset_file coco \ --coco_path ./data/coco \ --resume ./outputs/my_detr/checkpoint.pth \ --eval它会输出COCO风格的AP、AP50、AP75以及AP_small、AP_medium、AP_large。这些指标的直观意义:
- AP50:IoU阈值0.5下的平均精度,是检测界最常用的宽松指标。
- AP75:IoU阈值0.75下的平均精度,严苛一些,对框的定位精度更敏感。
- AP_small / AP_medium / AP_large:分别对小于32x32像素、32到96像素、大于96像素的目标统计AP。DETR在小物体上的AP通常低于大物体,这是Transformer系检测器的通病,后面会讲怎么缓解。
如果你发现AP_large不错、AP_small明显偏低,而且你的场景里小物体很多,那DETR原版可能不是最优选择,可以考虑后续看Deformable DETR或者DN-DETR,它们在小物体上有更优的表现。
5. 预测推理:让模型输出检测框
训练和评估跑通了,下一个环节是拿模型做推理,也就是输入一张图,得到检测框和类别。DETR的推理简单到让我有点感动。不用NMS,不用anchors,直接把图像预处理后喂给模型,拿输出做一次阈值过滤就行。
5.1 加载模型与权重文件
推理如果直接用官方代码里的main.py加--eval,只能输出评测指标,不会展示结果图。如果你要可视化检测效果,官方仓库其实没有提供现成脚本,需要自己写一段推理代码。
我参考官方demo逻辑,自己写过一套简化推理脚本。核心流程如下:
import torch from PIL import Image import torchvision.transforms as T import torch.nn.functional as F from models import build_model import argparse def get_args(): parser = argparse.ArgumentParser() parser.add_argument('--backbone', default='resnet50') parser.add_argument('--dilation', action='store_true') parser.add_argument('--num_classes', default=2) parser.add_argument('--num_queries', default=100) parser.add_argument('--position_embedding', default='sine') parser.add_argument('--checkpoint', required=True) parser.add_argument('--image_path', required=True) parser.add_argument('--threshold', default=0.7, type=float) return parser.parse_args() args = get_args() model, criterion, postprocessors = build_model(args) checkpoint = torch.load(args.checkpoint, map_location='cpu') model.load_state_dict(checkpoint['model']) model.eval()注意一点:加载checkpoint时,checkpoint['model']的键是模型权重字典,criterion是训练时用的丢失函数模块,推理时用不到,但build_model会返回它,所以接收一下即可。
5.2 图像预处理与后处理
DETR的推理图像预处理跟训练时略有不同。训练时会有随机裁剪缩放,推理时必须保持一致的分辨率处理逻辑,否则模型的性能会打折扣。
我的预处理代码如下:
transform = T.Compose([ T.Resize(800), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) img = Image.open(args.image_path).convert('RGB') img_tensor = transform(img).unsqueeze(0) with torch.no_grad(): outputs = model(img_tensor)outputs是一个dict,其中pred_logits的shape是[1, num_queries, num_classes + 1],因为背景也算一个类别;pred_boxes的shape是[1, num_queries, 4],输出的是归一化的[cx, cy, w, h]。
后处理逻辑:
probs = F.softmax(outputs['pred_logits'], dim=-1)[0] scores, labels = probs[..., :-1].max(dim=-1) keep = scores > args.threshold scores = scores[keep] labels = labels[keep] boxes = outputs['pred_boxes'][0][keep]这里有个小细节:pred_logits最后一维是背景类别,你要最大概率取在非背景的类别上,而不是在全部类别上取max,否则背景类别的得分可能干扰结果。这就是为什么上面代码用了probs[..., :-1]。
最后把归一化的框转换成像素坐标,并画到图上:
import cv2 import numpy as np img_np = np.array(img) h, w = img_np.shape[:2] boxes[:, 0::2] *= w boxes[:, 1::2] *= h for box, score, label in zip(boxes, scores, labels): x_c, y_c, bw, bh = box x1, y1 = int(x_c - bw / 2), int(y_c - bh / 2) x2, y2 = int(x_c + bw / 2), int(y_c + bh / 2) cv2.rectangle(img_np, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(img_np, f'{label.item()}:{score.item():.2f}', (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1) cv2.imwrite('result.jpg', img_np)跑出来如果你看到图片上框的位置很准,类别也正确,恭喜你,整个流程通了。如果框位置明显偏移,大概率是预处理时Resize的方式和训练不一致,或者反算坐标时宽高顺序写错了。
5.3 输出结果中的特殊Case
有时候你会遇到一个很有意思的情况:同一张图,同一个物体,模型输出了两个高置信度的重叠框。这在YOLO里基本算重复检测,要用NMS合并;但在DETR里,这种情况偶发也是正常的,因为DETR的查詢之间本身有交互,理论上不会重复大面积输出。如果真的出现了,多半是模型没有充分收敛,或者查询数量超过了对这张图有效物体的数量,可以适当调低阈值,或者增加训练轮数。
还有一种情况是,模型会找出训练集里没出现过的背景区域的框,置信度还挺高。那是因为模型把某个背景纹理误判成了目标类别,这在类别少、数据量小的场景下很常见。对策是:增加训练数据多样性,或者调高置信度阈值,而不是修改模型结构。
6. 常见问题与排查实录:我踩过的坑
6.1 训练时loss变成NaN
这个坑我印象最深。在自定义数据集上训练时,总loss在第15个epoch左右突然变成NaN,训练直接崩掉。
排查了一圈,最后定位到原因:数据集里有一张标注框宽高为负的图,DETR的训练逻辑在计算GIoU时,对负尺寸的框做了对数或者除法运算,直接产生了NaN。之前用YOLO训练,很多框架会默默处理掉这种非法标注,但DETR代码里并没有强校验。
解决方案很简单:在数据转换脚本里增加一个规则,丢弃宽高小于等于0的标注框,同时检查框是否超出图像边界。
if bw <= 0 or bh <= 0: continue if x1 < 0 or y1 < 0 or x2 > w or y2 > h: continue这个习惯建议从第一步就养成,别等训练崩了再回头查数据。
6.2 训练很久但mAP基本不动
另一个常见问题是:训练了二三十个epoch,loss看起来在下降,但验证集AP几乎没有上升。
我遇到过一次,最后发现是--num_queries设太小(比如10),而自己数据里单张图最多有20多个目标,查询数量不够,模型根本没有足够的输出槽位去预测所有物体。DETR的匹配机制是:真实框必须跟某个查询匹配,查询不够,部分物体就没有对应的查询能学到它,自然AP上不去。
解决办法是,统计自己训练集里单张图的最大目标数量,然后把这个数量乘以1.5到2,再设定为--num_queries。默认100对绝大多数场景是够用的,但你的场景如果特别密集,记得多加一点。
6.3 推理结果对不准,框是斜的
这里说的斜不是旋转框,而是预测框和物体实际轮廓明显不贴合。比如一辆车的检测框偏大或者偏小很多。
排查方向一般是两个。一是模型训练不充分,尤其是GIoU loss没有稳定收敛。你可以看看训练日志的loss_giou,如果它还在缓慢下降,说明模型还没学到位,继续训练。二是你自己的评估数据里,有些物体被遮挡严重,标注本身可能就不准,模型学到的框自然不准。这个问题没有特别好的办法,只能清理一下标注质量差的数据,或者考虑用更大更强的backbone。
6.4 常见问题速查表
| 问题现象 | 可能原因 | 解决思路 |
|---|---|---|
| loss变NaN | 标注框宽高为负或越界 | 校验数据格式,剔除非法标注 |
| AP长期很低 | 查询数量不足 | 统计最大目标数,增大num_queries |
| 推理框偏移 | 预处理不一致 | 统一训练/推理的resize与normalize参数 |
| 小物体AP极低 | DETR对小物体不敏感 | 提高输入分辨率,或换Deformable DETR |
| CUDA OOM | 显存不足 | 减小batch size、降低输入分辨率、开AMP |
| 加载checkpoint报错 | 模型类别数不匹配 | 修改--num_classes并重新加载对应权重 |
7. 最后的几点经验
跑通DETR只是第一步,真正有价值的是你理解了这套范式的运作逻辑。我自己在跑完DETR并跟YOLO做对比后,最大的感受是:DETR这类的端到端检测器,适合的场景是“检测流程越简单越好”和“模型结构需要自定义修改”的项目。如果你做的是工业化落地、追求极致的速度和绝对精度,YOLO系列目前的工程成熟度仍然更高。
但如果你做研究探索、做多模态任务、做检测和其他视觉任务的联合建模,DETR的基本盘非常扎实。从DETR出发,你可以很容易扩展到DAB-DETR、DN-DETR、Deformable DETR这些改进版本,每一步都可以在原有代码基础上小步改动,这种迭代效率是传统检测框架很难给的。
如果你打算在自己的数据集上认真使用DETR,我最后再分享一个实用技巧:不要一上来就训练完整模型。先用官方在COCO上的预训练权重,在你自己数据集上做迁移学习,只训练最后的检测头和Transformer解码器部分。通常只需要很少的轮数就能收敛,效果还好。具体做法是加载预训练权重后,把分类头的输出维度改成你自己的类别数,其他部分照搬,然后用小学习率微调整个网络,轮数10到20轮,效果往往比从头训练50轮还稳。
DETR这条路线,在目标检测领域已经算得上一股不可忽视的方向。我自己就是用这种方式快速验证了好几个基于修改查询表达方式的检测想法,省掉了大量调参时间。希望这篇内容能帮你少走弯路,早点跑通自己的DETR流程。