告别YOLO！基于Transformer的DETR实战：环境搭建+训练+预测全流程

在目标检测领域，YOLO、Faster R-CNN、SSD等基于卷积神经网络（CNN）的模型早已深入人心。它们依靠CNN提取特征，再通过预测头完成分类与边界框回归，但始终受限于CNN的局部特征提取能力。

而DETR（Detection Transformer）的出现打破了这一格局——它依然用CNN做特征提取，但核心部分引入了Transformer结构，通过编码器（Encoder）学习图像全局上下文，解码器（Decoder）完成预测，无需复杂的锚点设计，极大简化了目标检测流程。

今天就带来一篇全流程实战教程，教你用PyTorch框架，基于任意自定义数据集训练DETR模型，从环境搭建到最终的测试可视化！

一、三类目标检测模型

在动手之前，先快速理清目标检测的三大主流方向，帮你理解DETR的优势：

1. 单阶段检测器（Single-stage Detectors）

• 代表模型：YOLO、SSD、RetinaNet
• 核心逻辑：CNN骨干网络提取特征图后，直接在单网络中完成分类+边界框预测
• 优势：速度快，适合实时检测场景；劣势：小目标检测精度稍弱

2. 双阶段检测器（Two-stage Detectors）

• 代表模型：Faster R-CNN、R-CNN、Cascade R-CNN
• 核心逻辑：先通过区域提议（Region proposals）生成候选框，再基于多尺度特征图做精准预测
• 优势：小目标检测效果好；劣势：速度慢，不适合实时场景

3. 基于Transformer的检测器

• 代表模型：DETR、Deformable DETR、Conditional DETR
• 核心逻辑：CNN提取基础特征，Encoder学习图像全局关联（自注意力机制），Decoder转换特征完成预测
• 优势：无需锚点设计，全局特征捕捉能力强，适配自定义数据集；劣势：对GPU算力有一定要求

DETR作为Transformer类检测器的代表，兼顾了一定的速度与精度，且训练流程简洁，非常适合需要自定义数据集的场景。

二、实战步骤：从0训练DETR模型

1. 环境搭建：GPU+PyTorch是基础

DETR训练需要GPU支持，推荐两种方案：

• 本地搭建：需配置GPU支持的PyTorch环境（可参考官方教程）
• 云端方案：优先选择Kaggle（免费GPU训练时间更长）或Google Colab，无需本地配置

关键依赖安装（本地/云端通用）

1. 克隆DETR官方仓库

git clone https://github.com/facebookresearch/detr.git

2. 安装COCO数据集依赖

conda install cython scipypip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

3. numpy版本兼容处理（若报错）

pip install numpy==1.23

2. 数据集准备：COCO格式是关键

DETR默认支持COCO格式数据集，核心要求如下：

• 数据集来源：可使用Roboflow等平台下载（本文以猫狗检测数据集为例），也可自定义数据集
• 格式要求：必须严格遵循COCO格式（含图像文件+标注文件）
• 文件夹命名：无需修改代码时，需将文件夹命名为train2017（训练集）、val2017（验证集）、annotations（标注文件）

3. 模型训练

第一步：下载预训练权重

DETR提供多种预训练模型，本文选用常用的detr-r50-e632da11.pth，先创建权重文件夹并下载：

mkdir -p pretrained-weightswget https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth -O pretrained-weights/detr-r50-e632da11.pth

常用预训练模型参数参考（COCO 2017验证集效果）：

第二步：启动训练命令

激活环境后，修改coco_path为你的数据集路径，执行以下命令：

python main.py --dataset_file coco --coco_path /home/omer/vision-ws/detr-ws/dataset # 替换为你的数据集路径 --output_dir ./outputs # 模型输出路径 --batch_size 2 # 批次大小（根据GPU显存调整） --epochs 30 # 训练轮数 --lr_drop 40 # 学习率衰减轮数 --resume pretrained-weights/detr-r50-e632da11.pth # 预训练权重路径

• 注意：可根据需求调整batch_size（显存不足调小）、epochs（数据量大则增加）、lr_drop（学习率衰减时机）等参数
• 训练时长：取决于数据集大小和GPU性能，耐心等待即可

4. 模型预测：测试图像+可视化

训练完成后，outputs文件夹会生成模型权重文件，接下来用测试图像验证效果，全程复用训练环境。

第一步：导入依赖库

from PIL import Imageimport matplotlib.pyplot as plt%config InlineBackend.figure_format = 'retina'import torchfrom torch import nnfrom torchvision.models import resnet50import torchvision.transforms as Ttorch.set_grad_enabled(False);

第二步：定义常量与工具函数

自定义类别（以猫狗检测为例，可根据你的数据集修改）CLASSES = [‘N/A’, ‘cat’, ‘dog’]# 可视化颜色配置COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125], [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]# 图像预处理（标准化+尺寸调整）transform = T.Compose([ T.Resize(600), # 可根据需求调整尺寸 T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])# 边界框格式转换（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)# 边界框缩放至原图尺寸def rescale_bboxes(out_bbox, size): img_w, img_h = size b = box_cxcywh_to_xyxy(out_bbox) b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32) return b# 结果可视化函数def plot_results(pil_img, prob, boxes): plt.figure(figsize=(12,8)) plt.imshow(pil_img) ax = plt.gca() colors = COLORS * 100 for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors): ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=c, linewidth=3)) cl = p.argmax() text = f’{CLASSES[cl]}: {p[cl]:0.2f}’ ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor=‘yellow’, alpha=0.5)) plt.axis(‘off’) plt.show()

第三步：加载模型并预测

导入DETR模型构建模块import syssys.path.append(‘./detr’)from models.detr import build# 模型参数配置（与训练时一致）class Args: backbone = ‘resnet50’ dilation = False position_embedding = ‘sine’ num_classes = 2 # 猫狗两类（背景自动处理） hidden_dim = 256 nheads = 8 num_encoder_layers = 6 num_decoder_layers = 6 dim_feedforward = 2048 dropout = 0.1 enc_layers = 6 dec_layers = 6 pre_norm = False num_queries = 100 lr_backbone = 1e-5 set_cost_class = 1 set_cost_bbox = 5 set_cost_giou = 2 mask_loss_coef = 1 dice_loss_coef = 1 bbox_loss_coef = 5 giou_loss_coef = 2 eos_coef = 0.1 device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’ dataset_file = ‘coco’ masks = False aux_loss = True frozen_weights = Noneargs = Args()# 构建模型并加载训练权重model, criterion, postprocessors = build(args)checkpoint_path = ‘./detr/outputs/checkpoint.pth’ # 训练权重路径checkpoint = torch.load(checkpoint_path, map_location=args.device, weights_only=False)model.load_state_dict(checkpoint[‘model’], strict=False)model.to(args.device)model.eval() # 切换至评估模式# 加载测试图像并预测im = Image.open(‘testimage.jpg’) # 替换为你的测试图像路径img = transform(im).unsqueeze(0)img = img.to(args.device)# 模型推理outputs = model(img)# 筛选置信度≥0.5的预测结果probas = outputs[‘pred_logits’].softmax(-1)[0, :, :-1]keep = probas.max(-1).values > 0.5# 边界框缩放与可视化bboxes_scaled = rescale_bboxes(outputs[‘pred_boxes’][0, keep].cpu(), im.size)plot_results(im, probas[keep].cpu(), bboxes_scaled)

在这里插入图片描述

三、总结

DETR作为基于Transformer的目标检测模型，凭借简洁的架构和强大的全局特征捕捉能力，成为自定义数据集场景的优质选择。

读者福利：如果大家对大模型感兴趣，这套大模型学习资料一定对你有用

对于0基础小白入门：

如果你是零基础小白，想快速入门大模型是可以考虑的。

一方面是学习时间相对较短，学习内容更全面更集中。
二方面是可以根据这些资料规划好学习计划和方向。

作为一名老互联网人，看着AI越来越火，也总想为大家做点啥。干脆把我这几年整理的AI大模型干货全拿出来了。
包括入门指南、学习路径图、精选书籍、视频课，还有我录的一些实战讲解。全部免费，不搞虚的。
学习从来都是自己的事，我能做的就是帮你把路铺平一点。资料都放在下面了，有需要的直接拿，能用到多少就看你自己了。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以点击文章最下方的VX名片免费领取【保真100%】