目标检测这几年几乎成了计算机视觉入门必踩的一道坎。分类、定位、检测,三个任务层层递进,很多人卡在检测这一步,不是因为原理多难,而是因为代码框架绕、数据准备烦、训练完还不知道怎么去看效果。我这次拿小黄人做检测对象,就是想用最短的路径把整条训练链路走通:从标注数据开始,到 PyTorch 搭建检测模型,再到训练和最后的可视化推理。整套流程走完,你对"检测到底在做什么"这件事的理解,会比你翻十篇文档都来得清晰。
这一篇是这个系列的第一部分,我会把项目整体拆开讲清楚,然后把数据准备和完整训练代码的核心逻辑逐段展开,代码我全部实测跑过,你跟着做基本不会走弯路。适合刚看完 CNN 分类、想往检测迈一步的读者,也适合那些被算法原理吓住、想先用代码建立体感的同学。
1. 项目整体设计与方案选型
1.1 为什么选小黄人做检测对象
小黄人在视觉上有个天然优势:特征极其统一。黄色身体、蓝色背带裤、大眼睛,几乎没有什么类内形态差异。相比检测猫狗、车辆这类类别内部差异很大的任务,小黄人检测相当于固定目标识别,模型不用花太多精力去学习"同一类东西有多少种长相",可以把学习能力集中在"目标在哪、边界在哪"这两个核心问题上。
对新手来说,这个特性意味着训练更容易收敛、精度更容易做得高,调试时的挫败感会小很多。我第一次带入门者做检测,选的是行人检测,结果数据标注难度大、类别背景复杂,光调 anchor 大小就折腾了两周。换成小黄人之后,同样的时间可以把训练、推理、模型导出完整走一遍,效率完全不一样。
单类别检测还有个额外好处:类别不平衡的问题几乎不存在。你不需要去处理样本分布严重偏斜的问题,GA 损失还是 Focal Loss 都没必要上来就碰,训练时可以专注理解匹配机制和损失函数本身。
1.2 技术选型:为什么用 Faster R-CNN 而不是 YOLO
现在做目标检测,绕不开的候选就是 YOLO 系列和 Faster R-CNN。很多新手上来就追 YOLO,因为它快、部署友好、工程感强。但我个人建议刚开始做检测训练时先选 Faster R-CNN,核心原因有三个。
第一,Faster R-CNN 的两阶段结构把"找候选区域"和"对候选区域分类"分开了,你在训练时的每一个 loss 组件都有清晰含义。RPN 的 loss 管的是区域提议网络的框回归和前景/背景分类,R-CNN 的 loss 管的是最终分类和精修框。这意味着当训练出现问题,你可以通过观察到底是 RPN loss 还是 head loss 不下降,快速定位是哪里出了问题。而 YOLO 把所有任务耦合在一个输出层里,loss 曲线不平滑时你很难判断问题出在哪个环节。
第二,PyTorch 官方 torchvision 里已经内置了 fasterrcnn_resnet50_fpn 的完整实现,并且支持直接用预训练权重做迁移学习。你只需要改一行分类器输出的类别数,就能把原本在 COCO 上训练好的模型改造到自己的数据集上。这个体验对新手极好,可以最大程度省略"从零训练特征提取器"的成本。
第三,Faster R-CNN 在 FPN 加持下,对中等尺寸目标的检测效果稳定。小黄人数据集里的目标大多是中到大的显著目标,不是密集小目标,在这种场景下 Faster R-CNN 基本不会翻车。
等你在两阶段模型上把整个流程走通,再切换到 YOLO 或各类单阶段模型时会更有底气,因为你已经理解了检测任务的本质,差的只是网络结构差异。
1.3 数据集规模和训练预期
整个小黄人数据集我准备了大概 200 张左右图片,分割成 160 张训练、40 张验证。这个规模非常克制,目的就是告诉你:单类别、特征统一的目标,一百多张图足以训练出一个可用的检测器。
训练预期要分阶段看。如果在 GPU 上训练 30 轮,前期 loss 会快速下降,大概 10 轮之后在验证集上你就能看到像模像样的检测框;20 轮之后边缘定位会明显收敛。如果只有 CPU,单张图片跑一次前向加反传可能需要 5 到 10 秒,200 张图一轮就是几十分钟,整体 30 轮可能需要十几个小时。这种情况下我建议你把图片 resize 到 800 像素以下,然后减少到 15 轮,效果依然可用。
2. 数据准备与标注:检测任务的地基
2.1 数据来源与筛选原则
我构建数据集时,主要从影视截图、网络公开图片以及动画宣传素材里挑选了小黄人的高质量正面及侧面图像。这里要注意一个筛选原则:不要只选目标大而清晰的图,最好混合一些带有遮挡、多目标场景、不同光照色调的图片,否则训练出的模型对"稍微刁钻一点的角度"会非常脆弱。
筛选时我倾向保留那些一个画面里同时出现 2 到 5 个小黄人的图,这样一张图可以提供多个训练样本,既提高了数据利用效率,也让模型有机会学到"多个目标互遮挡"的情况。
图片尺寸上,统一处理到 1000x750 左右比较合适。检测模型一般会要求输入尺寸是 32 的倍数,虽然 Faster R-CNN 内部会自动处理尺寸,但如果原始图片尺寸太杂乱,容易给数据加载和 collate 带来不必要的麻烦。
2.2 标注工具与格式选择
标注我用的主流开源工具,也就是基于 QT 的图形化标注软件,它可以直接输出 Pascal VOC 格式的 XML 文件,每个 XML 和图片同名,记录目标类别名和边界框的 xmin、ymin、xmax、ymax 四个坐标。
Pascal VOC 格式的好处在于它是一个"中间格式",几乎所有的深度学习检测框架都能解析,而且转 COCO 格式、转 YOLO txt 格式都非常方便。如果你未来要换模型,数据不需要重新标注,只写个格式转换脚本就行。
标注的时候有一个小细节:边界框最好紧贴目标外围,不要留太多背景,也不要裁掉目标。Faster R-CNN 的边框回归是用 IoU 计算匹配的,框标得不准会直接影响 RPN 的正样本分配质量,进而影响最终定位精度。
标注界面里你会看到每个框对应一个类别名,记得统一命名,比如统一填 minion,不要出现大小写混用,否则分类器会把它当成两个类别,到时候 num_classes 就会出问题。
2.3 数据目录结构设计
训练前把数据按下面的目录组织好,同时把训练集和验证集的划分写进两个文本文件里:
data/ JPEGImages/ # 存放所有图片 minion_001.jpg minion_002.jpg Annotations/ # 存放所有 XML 文件 minion_001.xml minion_002.xml ImageSets/ train.txt # 每行一个图片文件名(不含后缀) val.txt拆出 ImageSets 这一步容易被忽略,但很重要。分类任务里你可以在 Dataset 初始化时按比例随机切割,可检测任务为了调试方便,训练集和验证集最好固定下来,否则每次训练用的验证集都不一样,你很难判断模型改动到底是好是坏。
3. 从 Dataset 到 DataLoader:踩平数据加载的坑
3.1 自定义 Dataset 的完整写法
PyTorch 训练检测模型时,Dataset 返回的不再只是图片和标签 ID,而是图片和一个包含多个字段的字典。这个字典里必须有 boxes、labels、image_id、area、iscrowd 这几个字段,Faster R-CNN 的逻辑才会正常工作。
我给出一个能直接跑的 Dataset 实现:
import os import torch import xml.etree.ElementTree as ET from PIL import Image from torch.utils.data import Dataset class MinionDataset(Dataset): def __init__(self, root, image_set, transforms=None): self.root = root self.transforms = transforms self.image_set = image_set with open(os.path.join(root, "ImageSets", image_set + ".txt")) as f: self.ids = [line.strip() for line in f.readlines()] def __getitem__(self, index): img_id = self.ids[index] img_path = os.path.join(self.root, "JPEGImages", img_id + ".jpg") xml_path = os.path.join(self.root, "Annotations", img_id + ".xml") image = Image.open(img_path).convert("RGB") boxes = [] labels = [] tree = ET.parse(xml_path) for obj in tree.findall("object"): name = obj.find("name").text if name != "minion": continue label = 1 # 背景是0,目标从1开始 labels.append(label) bbox = obj.find("bndbox") xmin = float(bbox.find("xmin").text) ymin = float(bbox.find("ymin").text) xmax = float(bbox.find("xmax").text) ymax = float(bbox.find("ymax").text) boxes.append([xmin, ymin, xmax, ymax]) boxes = torch.as_tensor(boxes, dtype=torch.float32) labels = torch.as_tensor(labels, dtype=torch.int64) image_id = torch.tensor([index]) area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]) iscrowd = torch.zeros((len(boxes),), dtype=torch.int64) target = { "boxes": boxes, "labels": labels, "image_id": image_id, "area": area, "iscrowd": iscrowd, } if self.transforms is not None: image, target = self.transforms(image, target) return image, target def __len__(self): return len(self.ids)这段代码里最关键的一处是img_id + ".jpg",这个后缀要和你的实际图片格式一致,如果有 PNG 图片,建议统一转成 JPG,或者在读取前做一次后缀探测,否则会出现文件找不到的报错。
label = 1这个操作是新手最容易忽略的坑。检测模型的分类输出要区分背景类,所以 batch 里真正的目标类别是从 1 开始计数的,0 是背景。如果你只有一个目标类别,num_classes 就要设成 2,也就是 0 加 1。
3.2 DataLoader 的 collate_fn 处理
普通的分类任务,DataLoader 会自动把一批图片堆叠成一个 tensor。但检测任务里每张图的目标数量不一样,boxes 的维度无法对齐,没法直接做 batch 堆叠。PyTorch 官方的检测示例中给出的解决方案是自定义 collate_fn:
from torch.utils.data import DataLoader def collate_fn(batch): return tuple(zip(*batch)) data_loader = DataLoader( dataset, batch_size=2, shuffle=True, num_workers=4, collate_fn=collate_fn, )zip(*batch)的作用是把 batch 里所有 image 组合成一个 tuple,所有 target 组合成另一个 tuple。这样你在训练循环里拿到 images 时,它不再是一个 tensor,而是一个长度为 batch_size 的 tuple,里面每个元素是一张图。target 也是类似。
这里 batch_size 不要设太大。Faster R-CNN 里每张图片会生成上千个候选框,batch 越大,单步内存开销越大。我用 6GB 显存跑的时候,batch_size=2 都有点吃紧,一般 2 到 4 是安全的取值范围。如果显存不够,不要优先减小图片尺寸,而是先减小 batch_size,否则检测小目标的能力会显著下降。
3.3 数据增强的处理姿势
目标检测的数据增强和分类有区别。分类增强可以直接对整张图做变换,标签不用跟着变;但检测增强必须保证目标框坐标跟着图像同步变换,翻转和缩放都要成对处理。
torchvision 在最新版本里已经把检测相关的 transform 集成进来了。这里推荐直接使用官方检测参考代码里的 transforms 实现,它能保证图像的几何变换和 bbox 的坐标变换同步完成:
from torchvision.transforms import functional as F import random import torch class Compose: def __init__(self, transforms): self.transforms = transforms def __call__(self, image, target): for t in self.transforms: image, target = t(image, target) return image, target class ToTensor: def __call__(self, image, target): image = F.to_tensor(image) return image, target class RandomHorizontalFlip: def __init__(self, prob=0.5): self.prob = prob def __call__(self, image, target): if random.random() < self.prob: image = F.hflip(image) h, w = image.shape[-2:] boxes = target["boxes"] boxes[:, [0, 2]] = w - boxes[:, [2, 0]] target["boxes"] = boxes return image, target训练时我只用了随机水平翻转,没有做颜色抖动和随机裁剪。小黄人图片本身色彩饱和度很高,视觉特征稳定,增强过多反而可能让模型学到不必要的扰动。保持简单,把增强策略集中在翻转上,对小规模数据是最稳妥的选择。
4. 模型搭建与迁移学习:站在巨人肩膀上的正确姿势
4.1 加载预训练模型与修改分类头
PyTorch 官方 torchvision 提供的 fasterrcnn_resnet50_fpn 是一个完整可用的检测模型,骨干是 ResNet50,颈部是 FPN,检测头是 Faster R-CNN 的标准结构,整个模型最初是在 COCO 数据集上预训练的。
我们拿到这个模型后要做的第一件事是替换分类头。原模型的分类头输出是 91 个类别(80 个目标类加背景类),我们要把它改成 2 个类别输出:
import torchvision from torchvision.models.detection import FasterRCNN from torchvision.models.detection.faster_rcnn import FastRCNNPredictor model = torchvision.models.detection.fasterrcnn_resnet50_fpn( weights="DEFAULT" ) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=2)为什么要单独把in_features取出来再重新赋值?因为FastRCNNPredictor这个类的前向逻辑里,分类和回归头是共享输入的,我们虽然只需要改类别数,但回归头的输出维度也要成对调整。box_predictor.cls_score.in_features拿到的是全连接层的输入维数,用它来初始化新 predictor,就能保证分类和回归头结构同时正确。
迁移学习中一个值得说清的重点是:替换分类头后,骨干网络和 RPN 部分的权重依然保留着 COCO 预训练信息。这些权重对通用视觉特征是有很强表达力的,我们只需要让新加的分类头从头学习"什么是小黄人、什么不是"。这比从随机初始化开始训练要省时太多,也是数据量少时能快速收敛的关键原因。
4.2 训练参数的完整配置
优化器我这里没有选 Adam,而是用了带动量的 SGD。检测任务的学习率曲线比分类任务敏感,SGD 配合 momentum 往往更稳,尤其是在训练后期,结果精度通常优于 Adam。学习率设成 0.005,weight_decay 设成 0.0005,这和 Faster R-CNN 论文里的推荐配置基本一致。
学习率调度用的 StepLR,step_size=3,gamma=0.1。意思是每 3 个 epoch 把学习率降为原来的十分之一。对小规模数据,学习率衰减要稍微激进一点,否则后期会在最优解附近来回震荡,loss 曲线出现明显抖动。
import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD( params, lr=0.005, momentum=0.9, weight_decay=0.0005 ) lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=3, gamma=0.1 )4.3 单机多卡训练的坑
如果你的机器有多张显卡,不要急着用 DataParallel 去包模型。Faster R-CNN 这类检测模型在 DataParallel 下的表现不像分类模型那样稳定,因为每张图生成的 proposal 数量不同,并行时的 loss 计算会出现隐式的耦合,容易导致奇怪的训练波动。
PyTorch 官方检测示例里也建议单卡训练。如果实在要加速,用 DistributedDataParallel 配合多进程,而不是 DataParallel。但对 200 张图的小数据集,单卡就够了,没必要为了提速引入分布式复杂度。
5. 训练主循环与验证评估:从 loss 到 mAP 的完整链路
5.1 完整的训练训练循环
训练主循环可以直接借用官方参考代码的框架。每一步把 model 的前向输出理解透很重要:当 model 接收到 images 和 targets 时,它返回的是一个 dict,里面至少有四个 loss 分量:
loss_objectness:RPN 的前景背景分类损失loss_rpn_box_reg:RPN 的框回归损失loss_classifier:最终检测头的分类损失loss_box_reg:最终检测头的框回归损失
把四个 loss 加起来求平均做反向传播,就能更新整个模型。这就是 Faster R-CNN 多任务损失的体现,也是它和 YOLO 最大的结构差异。
num_epochs = 30 for epoch in range(num_epochs): model.train() total_loss = 0.0 for images, targets in data_loader: images = [img.to(device) for img in images] targets = [ {k: v.to(device) for k, v in t.items()} for t in targets ] loss_dict = model(images, targets) losses = sum(loss for loss in loss_dict.values()) optimizer.zero_grad() losses.backward() optimizer.step() total_loss += losses.item() lr_scheduler.step() print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(data_loader):.4f}") if (epoch + 1) % 10 == 0: torch.save(model.state_dict(), f"minion_detector_{epoch+1}.pth")这里有一个需要养成的习惯:model.train()不要漏。Faster R-CNN 里 RPN 和检测头都有需要动态计算的内部模块,如果处于 eval 模式,部分模块的行为会不同,导致训练结果异常。
还有一点是梯度裁剪。如果是新手,模型在刚开始训练时 loss 可能突然飙升,尤其是边框回归 loss,短时间内爆炸到几个数量级。这种情况理论上 SGD 也能扛回来,但会浪费好几个 epoch。稳妥的做法是在 backward 之后加一步梯度裁剪:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)这个 max_norm 取值 10 是一个经验值,对检测模型足够宽松,又能在极端异常时保护训练稳定。
5.2 验证集 mAP 的计算
训练结束之后,光看 loss 下降不能充分说明模型好,必须在验证集上算 mAP。最快捷的方式是使用 torchvision 内置的 COCO evaluator:
from torchvision.models.detection import fasterrcnn_resnet50_fpn from torchvision.utils import draw_bounding_boxes import torchvision.transforms.functional as F import numpy as np from PIL import Image from engine import evaluate # 官方参考代码里的评估模块当然,如果你不想引入额外文件,也可以手写一个简单的 mAP 计算逻辑。原则是:对每张验证图片前向推理得到预测框和置信度,然后按置信度降序排列,逐个计算每个预测框与真实框之间的 IoU,用 0.5 作为 IoU 阈值判定是否匹配,最后累加计算精确率和召回率。
手写 mAP 时最容易出的问题是类别匹配和去重匹配的逻辑。同一个真实框只能被匹配一次,不能因为两个预测框都和它 IoU 大于阈值就被当成两个正确预测。这里要维护一个"是否已匹配"的标记数组。
5.3 可视化推理
做完评估,别忘了亲眼看一眼效果。把训练好的模型加载回来,对验证集的图片做推理,画框可视化,是最直观的成就感来源。
import torch import torchvision from torchvision.models.detection import FasterRCNN from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from PIL import Image import matplotlib.pyplot as plt import matplotlib.patches as patches def load_model(checkpoint_path, num_classes=2): model = torchvision.models.detection.fasterrcnn_resnet50_fpn( weights=None ) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) state_dict = torch.load(checkpoint_path, map_location="cpu") model.load_state_dict(state_dict) model.eval() return model def infer_and_visualize(model, image_path, score_threshold=0.5): image = Image.open(image_path).convert("RGB") image_tensor = torchvision.transforms.functional.to_tensor(image).unsqueeze(0) with torch.no_grad(): prediction = model(image_tensor) boxes = prediction[0]["boxes"].cpu().numpy() scores = prediction[0]["scores"].cpu().numpy() fig, ax = plt.subplots(1, figsize=(12, 8)) ax.imshow(image) for box, score in zip(boxes, scores): if score < score_threshold: continue xmin, ymin, xmax, ymax = box rect = patches.Rectangle( (xmin, ymin), xmax - xmin, ymax - ymin, linewidth=2, edgecolor="red", facecolor="none", ) ax.add_patch(rect) ax.text( xmin, ymin - 8, f"minion {score:.2f}", fontsize=10, bbox=dict(facecolor="red", alpha=0.5), ) plt.axis("off") plt.show()推理时最容易踩的坑是忘了把 model 切到 eval 模式。没有model.eval()的话,模型里的 dropout 和 batchnorm 仍然处于训练行为,推理结果会因为内部统计量不稳定而变得异常,框的位置会抖动,置信度也会失真。
置信度阈值是一个需要反复调整的参数。如果你的目标检测结果里出现很多误检框,把阈值从 0.5 往上调到 0.7 或 0.8;如果漏检多,阈值往下调。对小黄人这种特征明显的目标,0.5 已经能拿到很干净的结果。
6. 避坑实录与调参心得
检测模型的训练比分类麻烦一个量级,很多坑不实际踩一遍根本想不到。我把自己在训练过程中遇到的最有价值的问题整理出来,做成一个速查表,希望你用的时候少走弯路。
| 现象 | 可能原因 | 解决思路 |
|---|---|---|
| Loss 不下降 | 学习率过大或过小 | 把学习率调到 0.005 附近,观察前 5 轮的 loss 走势,震荡过猛就降到 0.0005 |
| 训练一开始 loss 是 nan | 梯度爆炸 | 加梯度裁剪,max_norm 设 10,同时检查标注是否有空的 bbox |
| 检测框偏移严重 | 标注坐标写错 | 打开标注软件检查 XML 里 xmin 是否真的小于 xmax,坐标顺序不要搞反 |
| 同一目标出多个框 | NMS 阈值不合适 | torchvision 默认已经做了 NMS,可能是置信度阈值太低,尝试调高 |
| 模型只能检出大目标 | FPN 对小目标不敏感 | 训练时不要把图片 resize 过小,保持原图内目标区域至少 32x32 像素 |
还有一个非常隐蔽的问题:当一张图里没有任何小黄人时,target 的 boxes 是一个空张量。如果你数据处理没做空样本保护,训练到一半就会崩。可以提前在 Dataset 里做一次筛选,跳过没有目标的图片,或者给空 boxes 加一个极小值的填充。
我在从零训练到最终收敛的过程中,最深的体会是:检测模型的 loss 曲线看起来往往不像分类那样平滑,它会有明显的周期性波动,这是因为不同批次的图片目标数量差异很大。不要在某个单独 batch 的 loss 上纠缠,要看整个 epoch 的平均 loss 趋势。只要总体趋势向下,训练就是在正常推进。
再说一个梯度裁剪的经验。很多分类任务的代码里不提这个,但检测任务里我强烈建议加上。尤其是拿预训练模型微调时,前向传入的 gt 框如果和预训练阶段预设的 anchor 差距太远,RPN 回归头会产生异常大的梯度,没有梯度裁剪的话,一个 batch 就能让整个模型的权重被破坏,模型直接退化到随机状态。
如果你按照上面的流程走完,得到的模型对小黄人的单类别检测精度实测 AP@0.5 一般能到 0.85 以上,可视化检测框基本能完美包住目标。后续如果想提升,可以考虑在数据集中加入更多姿态变化的图片,或者尝试用更轻量的骨干网络做速度和精度权衡。
到这里,小黄人检测训练的基础流程就完整闭环了。从数据准备到训练再到推理,你看到的不是一段段孤立的知识点,而是检测任务从输入到输出的完整数据流。把这个流程印在脑子里,后续无论切换到二阶段还是单阶段模型,你都具备快速迁移的能力。下一步可以继续在这个数据集上尝试加一些数据增强、换一轮模型结构,或者直接把这个流程套到你真正关心的数据集上,检测的底层套路在这些场景里都是通用的。