资讯动态

小黄人目标检测实战:PyTorch实现Faster R-CNN训练全流程

发布时间:2026/10/13 22:20:29 来源:尧图企业网站定制
目标检测这几年几乎成了计算机视觉入门必踩的一道坎。分类、定位、检测三个任务层层递进很多人卡在检测这一步不是因为原理多难而是因为代码框架绕、数据准备烦、训练完还不知道怎么去看效果。我这次拿小黄人做检测对象就是想用最短的路径把整条训练链路走通从标注数据开始到 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, transformsNone): 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, dtypetorch.float32) labels torch.as_tensor(labels, dtypetorch.int64) image_id torch.tensor([index]) area (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]) iscrowd torch.zeros((len(boxes),), dtypetorch.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_fnfrom torch.utils.data import DataLoader def collate_fn(batch): return tuple(zip(*batch)) data_loader DataLoader( dataset, batch_size2, shuffleTrue, num_workers4, collate_fncollate_fn, )zip(*batch)的作用是把 batch 里所有 image 组合成一个 tuple所有 target 组合成另一个 tuple。这样你在训练循环里拿到 images 时它不再是一个 tensor而是一个长度为 batch_size 的 tuple里面每个元素是一张图。target 也是类似。这里 batch_size 不要设太大。Faster R-CNN 里每张图片会生成上千个候选框batch 越大单步内存开销越大。我用 6GB 显存跑的时候batch_size2 都有点吃紧一般 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, prob0.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( weightsDEFAULT ) in_features model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor FastRCNNPredictor(in_features, num_classes2)为什么要单独把in_features取出来再重新赋值因为FastRCNNPredictor这个类的前向逻辑里分类和回归头是共享输入的我们虽然只需要改类别数但回归头的输出维度也要成对调整。box_predictor.cls_score.in_features拿到的是全连接层的输入维数用它来初始化新 predictor就能保证分类和回归头结构同时正确。迁移学习中一个值得说清的重点是替换分类头后骨干网络和 RPN 部分的权重依然保留着 COCO 预训练信息。这些权重对通用视觉特征是有很强表达力的我们只需要让新加的分类头从头学习什么是小黄人、什么不是。这比从随机初始化开始训练要省时太多也是数据量少时能快速收敛的关键原因。4.2 训练参数的完整配置优化器我这里没有选 Adam而是用了带动量的 SGD。检测任务的学习率曲线比分类任务敏感SGD 配合 momentum 往往更稳尤其是在训练后期结果精度通常优于 Adam。学习率设成 0.005weight_decay 设成 0.0005这和 Faster R-CNN 论文里的推荐配置基本一致。学习率调度用的 StepLRstep_size3gamma0.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, lr0.005, momentum0.9, weight_decay0.0005 ) lr_scheduler torch.optim.lr_scheduler.StepLR( optimizer, step_size3, gamma0.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_objectnessRPN 的前景背景分类损失loss_rpn_box_regRPN 的框回归损失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(fEpoch {epoch1}/{num_epochs}, Loss: {total_loss / len(data_loader):.4f}) if (epoch 1) % 10 0: torch.save(model.state_dict(), fminion_detector_{epoch1}.pth)这里有一个需要养成的习惯model.train()不要漏。Faster R-CNN 里 RPN 和检测头都有需要动态计算的内部模块如果处于 eval 模式部分模块的行为会不同导致训练结果异常。还有一点是梯度裁剪。如果是新手模型在刚开始训练时 loss 可能突然飙升尤其是边框回归 loss短时间内爆炸到几个数量级。这种情况理论上 SGD 也能扛回来但会浪费好几个 epoch。稳妥的做法是在 backward 之后加一步梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm10.0)这个 max_norm 取值 10 是一个经验值对检测模型足够宽松又能在极端异常时保护训练稳定。5.2 验证集 mAP 的计算训练结束之后光看 loss 下降不能充分说明模型好必须在验证集上算 mAP。最快捷的方式是使用 torchvision 内置的 COCO evaluatorfrom 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_classes2): model torchvision.models.detection.fasterrcnn_resnet50_fpn( weightsNone ) 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_locationcpu) model.load_state_dict(state_dict) model.eval() return model def infer_and_visualize(model, image_path, score_threshold0.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, linewidth2, edgecolorred, facecolornone, ) ax.add_patch(rect) ax.text( xmin, ymin - 8, fminion {score:.2f}, fontsize10, bboxdict(facecolorred, alpha0.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 就能让整个模型的权重被破坏模型直接退化到随机状态。如果你按照上面的流程走完得到的模型对小黄人的单类别检测精度实测 AP0.5 一般能到 0.85 以上可视化检测框基本能完美包住目标。后续如果想提升可以考虑在数据集中加入更多姿态变化的图片或者尝试用更轻量的骨干网络做速度和精度权衡。到这里小黄人检测训练的基础流程就完整闭环了。从数据准备到训练再到推理你看到的不是一段段孤立的知识点而是检测任务从输入到输出的完整数据流。把这个流程印在脑子里后续无论切换到二阶段还是单阶段模型你都具备快速迁移的能力。下一步可以继续在这个数据集上尝试加一些数据增强、换一轮模型结构或者直接把这个流程套到你真正关心的数据集上检测的底层套路在这些场景里都是通用的。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑