PyTorch目标检测入门:用Faster R-CNN训练小黄人检测模型
如果你也搜过“PyTorch 目标检测怎么入门”大概率和我一开始一样面对一堆模型名称、配置参数、官方教程里跳来跳去的术语半天不知道从哪里下手。分类任务一抓一大把教程但检测任务要同时输出“位置”和“类别”代码链路长踩坑点多没有个像样的练手项目根本坚持不下来。所以我干脆选了个最顺手的对象——小黄人写了一个完整的检测训练流程这是系列第一篇重点说清楚数据集、数据加载、模型搭建和训练全流程。为什么选小黄人而不是猫狗等下会细讲。先给你一个结论目标检测第一课对象越“好标”越好而小黄人这种颜色鲜明、外形完整、经常成群出现、姿态还挺丰富的卡通形象比真实世界里的猫猫狗狗更适合当第一个项目。这篇文章是给有一定PyTorch基础、但还没碰过检测任务的人看的我会把每一步为什么这么写、有哪些容易栽进去的坑都讲清楚你照着跑一遍能拥有一个真正可以检测小黄人的模型。1. 目标检测开题为什么“小黄人”是最顺手的练手对象我第一次做目标检测的时候上来就套了个通用模型去跑 COCO 数据集结果光看文档就花了两天训练出来的效果还说不清楚是数据问题还是代码问题。后来换了个思路把任务规模压缩到“只检测一种目标”先把链路跑通再谈优化。小黄人就是这个思路下最合适的对象。1.1 类别少但空间形态丰富目标检测模型要学两件事一是“这是什么”二是“这东西在哪、占多大”。如果只用一个类别分类那部分压力会小很多但位置和尺寸的多样性依然能充分暴露问题。小黄人天然具备这个优势。单人出现在画面中央的多人聚在一起互相遮挡的远处只有指甲盖大小的近到只截到半个身体的还有各种歪头、举手、奔跑姿势。你在标注的时候就会发现同一张图里经常有好几个目标而且大小差异很明显。这种数据分布非常适合让模型学会“框得准”也方便后续验证 NMS非极大值抑制到底做了什么——有时候两个框重叠得很厉害模型怎么决定留下哪一个看一眼输出就知道了。颜色也是加分项。小黄人的标志性黄色和深蓝工装在大多数背景里都比较突出背景分割相对容易模型不需要花太多精力去学“这是背景纹理还是目标的一部分”。这一点对新手阶段非常重要因为目标检测里“误检”是最让人头疼的问题之一选一个目标显著的数据集能把这类问题压到最低。1.2 数据规模与标注成本可控目标检测的数据集不像分类那么好凑。分类任务一张图一个标签网上爬图很容易而检测要人工画框一张图几十秒到几分钟不等。小黄人这个选题的另一个好处就是数据量不需要大。我当时收了两百多张图剔除掉重复和质量差的最后留了 150 张左右标注一下午就搞定。对比一下如果是做猫狗检测真实环境里猫的姿态差异极大背景杂乱遮挡严重没有几百张高质量标注图很难收敛小黄人因为是卡通角色画风相对统一150 张训练图配 30 张验证图已经能训出一个肉眼可用的模型。这种“数据规模小但任务完整”的项目最适合第一课。它能让你把所有环节都过一遍而不会被数据准备淹没。1.3 这个系列我会怎么安排这个系列我计划至少写三篇。第一篇是今天的重点搭建完整训练流程包括数据准备、自定义 Dataset、Faster R-CNN 迁移学习、训练循环、Loss 观察和初步推理。第二篇会补上评测部分讲 mAP 怎么算、到底怎么判断模型有没有变好以及针对漏检和误检的调参思路。第三篇可能会讲模型导出和部署把训练好的模型转成可落地的格式。所以今天这篇文章里我不会把所有细节都铺开而是尽量把“从原始图片到能跑通训练的完整链路”讲透。你跟着跑完这一篇手里的就是一个能接收图片并输出检测框的模型后面再想优化就有清晰的抓手了。2. 训练集准备从收集图片到XML标注很多人学目标检测花最多时间的反而不是看论文而是“怎么把图片变成模型能吃的数据”。这一步没做好后面全是灾难。我见过不少人训练到一半发现数据集里混了没有标注的文件、标注框坐标越界、标签名字不统一……所以我建议你先沉下心把数据集准备这一步走扎实。2.1 图片收集与初步筛选图片来源没那么多讲究我当时就是从公开渠道收集小黄人的壁纸、剧照、表情包合集再去掉带水印或者截得特别模糊的。收集的时候先别管“好不好标”原始图片数量越多越好因为后面筛完会缩水很多。收集完之后要过一遍初筛重点看三个问题目标太小有些全景图里小黄人只占几十个像素这种新手阶段直接删掉。目标太小会导致标注框只有二三十像素模型很难学Loss 也容易不稳定。严重遮挡比如两个小黄人叠在一起几乎看不出边界或者被文字、道具挡住大半。留着这种图你自己标注时都拿不准框线该画到哪模型当然更学不会。大尺寸重复图同一种构图的高清图经常出现很多张选一张代表即可避免训练集过于重复导致过拟合。初筛完我通常会看一眼整体分布单人图、多人图、不同背景的图比例是否均衡。我建议至少保留三成左右的多人图因为目标检测模型如果只在单目标图片上训练验证时一出现多个目标就漏检严重。2.2 LabelImg手工标注流程标注工具我推荐 LabelImg它足够简单安装也不费劲。先明确一点标注格式选 Pascal VOC也就是每张图片对应一个 XML 文件框的坐标以左上角为原点、以像素为单位记录。装好之后的工作流大概是这样的打开 LabelImg左侧选择“Open Dir”指定存放图片的文件夹。再选择“Change Save Dir”指定 XML 输出目录。按键盘 W 开始画框在小黄人身上拉一个矩形尽量贴住目标的边缘不要留太多背景。弹出的对话框里输入类别名统一叫小黄人的英文拼写比如 minion。画完一张图按 CtrlS 保存按 D 翻到下一张。听起来很简单但实际操作里有几个容易出问题的点。第一类别命名必须全程一致。不要这张图写 minion那张图写 Minion大小写不同在 XML 里就是两个类别模型训练时类别 id 会错乱。第二每张图都要检查有没有漏标。我习惯标完一批后把图片和 XML 一起用脚本抽出来画一遍框肉眼确认没有遗漏。第三像小黄人这种卡通角色如果一张图里只有侧脸或者只有上半身也照样要标不要因为它“不完整”就跳过模型需要学习各种截断情况。另外我建议这个阶段把所有图都放进同一个文件夹不要按主题分子目录。不是不能分而是对第一版项目来说没必要后面反而会让路径处理多出很多分支。2.3 数据划分与目录组织训练集和验证集划分我直接从全部图片里随机抽 8:2注意划分前先把文件全部拷贝到一个干净目录里不要动原始素材。最终目录结构可以长这样minion_detection/ ├── data/ │ ├── annotations/ # 所有 XML 标注 │ ├── images/ # 所有图片 │ ├── train.txt # 训练集图片文件名每行一个 │ ├── val.txt # 验证集图片文件名每行一个 │ └── label_map.json # 类别映射train.txt 里存的是不带扩展名的文件名后面对应到 image、xml 都很方便。label_map.json 我直接手写{ minion: 1 }这里有个新手特别容易犯的错误把类别 id 从 0 开始。但 PyTorch 检测模型的类别 0 是固定背景类如果小黄人标成 0模型会把它当成背景直接忽略Loss 还会出现奇怪的行为。所以要记住背景占 0第一个目标类别从 1 开始。2.4 用脚本检查标注文件手动标注难免出错强烈建议在训练之前写一个几十行的检查脚本把所有 XML 全部读一遍自动排查这几类问题文件名对应的图片是否存在XML 里的 width、height 是否和实际图片分辨率一致xmin、ymin、xmax、ymax 是否都在图片尺寸范围内且 xmin xmax、ymin ymax有没有空标注文件一个object都没有类别名是否都在 label_map 里。这一步我吃过亏。有一次我批量处理图片把某几张转成了灰度图XML 里记录的还是旧的路径训练到一半突然报错回头查才发现是数据加载时文件找不到。提前跑一遍脚本能省一晚上的排查时间。3. 自定义DatasetPyTorch检测模型的输入格式数据准备好了接下来是把图片和标注转换成 PyTorch 能用的数据结构。这一步看着不复杂但检测任务的 Dataset 写法比分类任务麻烦不少而且很多人第一版就死在 target 字典格式不对上。3.1 target字典该填哪些字段常规分类任务里__getitem__返回的是(image, label)label 通常是一个整数。检测任务不一样一张图里有多个目标所以 label 要换成字典官方模型要求的字段如下字段类型含义boxesFloatTensor[N, 4]N 个目标的边界框坐标labelsInt64Tensor[N]每个目标的类别 idimage_idInt64Tensor[1]图片唯一编号areaFloatTensor[N]每个框的面积iscrowdUInt8Tensor[N]是否为难样本单目标检测填 0新手最容易忽略的是area和iscrowd。有些教程里只给 boxes 和 labels训练也能跑原因是模型内部计算面积时会从 boxes 推导iscrowd 也会默认处理。但如果后面你想用 COCO 风格的评测工具或者想用官方更完整的训练管线最好还是把所有字段都补齐。养成习惯后面少踩坑。3.2 XML解析与boxes加载这里给出一个朴素的 Dataset 实现我没有加花哨的数据增强先把最基础的版本讲清楚。import os import json 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, imageset_txt, label_map): self.root root self.image_dir os.path.join(root, images) self.annotation_dir os.path.join(root, annotations) self.label_map label_map with open(imageset_txt, r) as f: self.image_ids [line.strip() for line in f if line.strip()] def _read_annotation(self, image_id): xml_path os.path.join(self.annotation_dir, f{image_id}.xml) tree ET.parse(xml_path) root tree.getroot() boxes, labels, area, iscrowd [], [], [], [] for obj in root.iter(object): name obj.findtext(name) if name not in self.label_map: continue bndbox obj.find(bndbox) xmin float(bndbox.findtext(xmin)) ymin float(bndbox.findtext(ymin)) xmax float(bndbox.findtext(xmax)) ymax float(bndbox.findtext(ymax)) boxes.append([xmin, ymin, xmax, ymax]) labels.append(self.label_map[name]) area.append((xmax - xmin) * (ymax - ymin)) iscrowd.append(0) boxes_t torch.as_tensor(boxes, dtypetorch.float32) labels_t torch.as_tensor(labels, dtypetorch.int64) area_t torch.as_tensor(area, dtypetorch.float32) iscrowd_t torch.as_tensor(iscrowd, dtypetorch.uint8) image_id_t torch.tensor([int(image_id)], dtypetorch.int64) return { boxes: boxes_t, labels: labels_t, image_id: image_id_t, area: area_t, iscrowd: iscrowd_t, } def __getitem__(self, idx): image_id self.image_ids[idx] image Image.open(os.path.join(self.image_dir, f{image_id}.jpg)).convert(RGB) target self._read_annotation(image_id) return image, target def __len__(self): return len(self.image_ids)几个细节要解释一下。坐标用的确实是原始像素坐标这是因为我这版代码没有对图片做任何缩放保持原尺寸输入模型。你可能在一些教程里看到“坐标要归一化到 0~1”那通常是因为图片被统一 resize 了。对第一版来说保持原尺寸能省掉大量坐标同步的逻辑。另外我用的是 PIL 直接返回Image对象而不是torchvision.transforms.ToTensor()之后的结果因为模型内部会自己做标准化和类型转换。这里其实有讲究torchvision检测模型期望输入是一个 PIL 图像列表而非预处理好的张量列表这一点和分类任务的习惯不太一样。一开始我总习惯在外面把图片转成张量结果模型反而报错后来才明白原版流程是模型内部处理。3.3 collate_fn与“可变尺寸”包装自定义好 Dataset 之后如果你直接把它丢进DataLoader大概率会碰到一个报错提示无法把 list of dict 转换为 Tensor。原因是检测任务里每个 batch 的样本框数不一样默认的 collate 函数无法把它们堆叠。解决办法也很简单写一个自定义的 collate_fn把整个 batch 的图片和 target 原样收集成两个元组def collate_fn(batch): images [item[0] for item in batch] targets [item[1] for item in batch] return images, targets然后 DataLoader 里加上collate_fncollate_fn即可。需要注意batch_size在这个阶段不要开太大因为图片尺寸没统一显存占用波动会比较大2 到 4 是比较稳妥的选择。3.4 为什么不建议第一版上随机裁剪很多做图像分类的人在数据增强里习惯了随机裁剪、随机旋转、随机颜色抖动。但目标检测里凡是改变几何位置的操作都必须同步修改 boxes 坐标这一步非常容易出 bug。比如随机裁剪你裁剪之后框的坐标必须减去裁剪区域的左上角偏移还要把超出裁剪边界的框截断甚至丢弃完全出界的框。如果哪一步忘了写模型看到的框和真实目标对不上Loss 会变得非常诡异。所以我的建议是第一版只做水平翻转。水平翻转的坐标变换非常简单对一张宽为 W 的图翻转后xmin_new W - xmaxxmax_new W - xminymin、ymax 不变。如果你想加直接在 Dataset 里用随机数控制即可。如果连水平翻转都不想写那也行。第一版的目标是先跑通数据增强放到第二篇再系统加。4. 基于预训练Faster R-CNN两行代码换掉分类头模型部分我选了 PyTorch 官方实现的 Faster R-CNN。理由后面展开先说结论它是目前理解检测完整流程最合适的起点而且代码量极少可以让我们把精力放在数据链路和训练调试上。4.1 为什么不用YOLO/SSD作为第一课选 YOLO 的人通常是被它的速度和部署生态吸引的但 YOLO 系列版本多、配置项复杂光理解 anchor、stride、cfg 文件就要花不少时间。SSD 的思路虽然也不错但 PyTorch 官方仓库里它没有 Faster R-CNN 维护得那么积极。Faster R-CNN 是一个两阶段模型第一阶段 RPN 负责“提出候选框”第二阶段 Fast R-CNN 对候选框做分类和精修。这种拆解非常直观你训练时能先后看到 RPN 的损失和检测头的损失出了问题也知道往哪边排查。而且 torchvision 官方实现了 ResNet50 FPN Faster R-CNN 的完整结构预训练权重现成迁移学习只需改一个分类头。对新手来说最重要的是“先跑通再深入”。Faster R-CNN 在这个目标上完成度最高。4.2 模型构建代码直接看代码import torch from torchvision.models.detection import fasterrcnn_resnet50_fpn from torchvision.models.detection.faster_rcnn import FastRCNNPredictor def build_model(num_classes2): model fasterrcnn_resnet50_fpn(pretrained_backboneTrue) in_features model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor FastRCNNPredictor(in_features, num_classes) return model这里num_classes2是因为分类数 背景 小黄人 2。如果你以后想加一个“香蕉”类别就改成 3。FastRCNNPredictor是官方提供的检测头替换接口它会根据in_features自动构造新的分类器和回归器。回归器的输出维度是num_classes * 4因为每个类别都要回归一组框坐标偏移。注意我把pretrained_backbone设置成了True也就是说我们用 ImageNet 上训练好的 ResNet50 作为骨干网络。这种迁移学习方式能显著降低训练难度尤其在这种小规模数据集上效果比从头训练好得多。如果你只想用特征提取器的参数、不想保留检测头原来的参数那就是当前这种写法官方默认加载的是完整模型的权重直接替换掉最后的 box_predictor其余头全保留为 COCO 预训练状态这在实际使用中效果也很好。4.3 冻结策略怎么选有人会问既然只是微调要不要把 backbone 冻结只训练检测头我建议第一版全量微调也就是所有层都参与训练。原因很简单目标检测任务对位置特征的要求比分类高很多backbone 冻结后它只能输出 ImageNet 分类用的特征未必能很好地表达“小黄人的边界在哪里”。全量微调虽然训练慢一点但收敛过程平滑效果明显更好。如果你的显卡显存确实不够可以退一步只冻结 backbone 的前两三个 stage让后面的 stage 继续训练。但我不推荐在新手阶段折腾这种策略因为它会让问题变复杂。4.4 预训练权重下载失败怎么办torchvision 在首次运行时会自动下载预训练权重。国内网络环境有时候下载不稳定卡在进度条半天没动静这是正常现象。优先建议手动下载权重文件放到本机的缓存目录。不同版本放的位置略有差异但通常在~/.cache/torch/hub/checkpoints/。下载完成后文件名要和代码里预期的一致比如fasterrcnn_resnet50_fpn_coco-258fb6c6.pth。如果你本地已经有这个文件直接把pretrained_backboneTrue改成传入weights路径也是可以的。这里要说一个通用经验不要让程序启动时每次去远程下载大文件。第一次下载成功之后最好把这个权重文件备份到一个固定目录下次直接用本地路径加载省时间也省心。5. 训练循环与Loss观察模型搭好了数据链路也通顺了接下来就是实际上手训练。训练循环本身不难难的是你看不懂 Loss不知道模型到底有没有在学习。5.1 训练主循环与优化器配置这里我直接展示一个可运行的训练脚本骨架省略了参数解析和日志记录的部分保留核心逻辑。import torch from torch.optim import SGD from torch.optim.lr_scheduler import MultiStepLR from torch.utils.data import DataLoader def train_one_epoch(model, loader, optimizer, device): model.train() total_loss 0.0 for images, targets in 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()) loss_value losses.item() total_loss loss_value optimizer.zero_grad() losses.backward() optimizer.step() return total_loss / len(loader)主流程里要注意model(images, targets)的调用方式。训练模式下传入 targets 之后模型返回的是一个 dict包含四个子损失全部加总才是最终的优化目标。优化器和学习率调度的配置我参考了 PyTorch 官方检测参考代码的常见做法params [p for p in model.parameters() if p.requires_grad] optimizer SGD(params, lr0.005, momentum0.9, weight_decay0.0005) scheduler MultiStepLR(optimizer, milestones[8, 11], gamma0.1)大概训练 12 个 epoch。如果你数据集更小可以减到 8 个如果验证集上的效果还在提升继续训也行。关键不是固定 epoch 数而是观察 Loss 的走势。5.2 四个Loss分别代表什么训练阶段会输出四个子损失新手看了容易懵这里先总结成一张表Loss 名称属于哪部分在学什么loss_objectnessRPN候选框内是否有目标loss_rpn_box_regRPN候选框位置的精修loss_classifier检测头候选框里的目标属于哪一类loss_box_reg检测头最终框坐标的回归总 Loss 是四者之和。正常训练时四个损失应该都呈下降趋势其中 loss_objectness 和 loss_classifier 下降得会比较平滑两个回归损失波动大一些尤其在训练初期。这不是 bug回归任务本身就比分类难收敛。我特别注意观察 loss_classifier 的变化。因为我们的任务类别很少它如果一直不降说明模型可能根本没有把目标和背景区分开这时要回头检查数据而不是急着调学习率。5.3 Loss变NaN的排查顺序NaN 是我碰到最多的训练异常没有之一。它的坑点在于报错不一定出现在触发 NaN 的那一步可能训练好几轮之后才突然中断让你很难定位。我建议按下面这个顺序排查先检查数据。把 DataLoader 里随机抽的样本都打印出来看一眼 boxes 有没有 NaNxmin 是否小于 0xmax 是否大于图片宽度。坐标越界是最常见原因。再检查标签。确认没有空标注样本没有类别 id 为 0 的目标。如果你在标注时不小心把第一个类别写成了 0模型会把目标当背景Loss 容易爆炸。检查学习率。lr0.005 在 batch_size 较小时可能偏大如果数据规模只有一百多张建议降到 0.002 或 0.001 试一下。排查 loss_bbox。如果训练到一半 loss_box_reg 先变 NaN多半是回归目标里有异常框比如面积为零或宽高为负数。5.4 训练慢/显存不够的应急方案很多人用自己的笔记本来训练第一次跑就发现显存直接被占满。解决办法有几个按优先级排序减小 batch_size。对于这个数据集batch_size1 或 2 完全够用只是训练时间会变长。限制图片最大尺寸。fasterrcnn_resnet50_fpn默认会将图片限制到 1333x800 左右如果你的图片分辨率特别高可以自己写一个预处理函数把长边压到 800 以内。注意 boxes 也要按相同比例缩放。换更小的骨干网络。把 ResNet50 换成 ResNet18 或者 MobileNet 的检测版本显存占用能降很多但精度也会受影响。这个属于后面的进阶话题。我自己的经验是150 张训练图batch_size2普通显卡大概跑二十分钟就能出效果。如果你的机器是纯 CPU那跑一遍要很久建议先用一小部分数据验证代码能跑通再上全量。6. 推理验证与效果调整训练完成之后最激动人心的环节就是把自己拍的、或者网上随便找的小黄人图片喂给模型看它能不能正确框出来。第一次看到模型跑出结果的时候那种成就感确实是分类任务给不了的。6.1 推理与可视化代码推理时需要把模型切到 eval 模式关闭梯度计算然后直接传入一张图片。模型返回的是一个 list每个元素对应输入图片的检测结果内容是 dict包含 boxes、labels、scores。import torch from PIL import Image, ImageDraw def predict(model, image_path, device, score_threshold0.5): model.eval() image Image.open(image_path).convert(RGB) with torch.no_grad(): prediction model([image.to(device)])[0] boxes prediction[boxes].cpu() scores prediction[scores].cpu() keep scores score_threshold draw ImageDraw.Draw(image) for box, score in zip(boxes[keep], scores[keep]): x1, y1, x2, y2 box.int().tolist() draw.rectangle([x1, y1, x2, y2], outlinered, width3) draw.text((x1, y1 - 10), f{score:.2f}, fillred) image.show()这里稍微解释一下scores是什么。它表示每个框属于某个类别的置信度范围在 0 到 1 之间。阈值设得越低框越多误检也越多设得越高漏检越多。0.5 是一个比较平衡的起点。另外你可能会好奇模型输出的框有没有经过 NMS实际上 torchvision 官方实现里已经在推理阶段默认做过一次置信度阈值过滤和 NMS输出结果中每个目标通常是“已经压掉重复框”之后的结果。新手阶段你不需要自己再写 NMS但应该知道它已经发生了。6.2 漏检/误检的调整思路跑完推理后你几乎肯定会遇到不满意的结果比如该框出来的没框出来或者背景被误检成了目标。这里给几个立即可用的调整思路。漏检多先看置信度阈值。把 score_threshold 从 0.5 降到 0.3如果框出来了说明模型本身学到了只是置信度偏低可以通过多训练几个 epoch 或者增强数据来解决。如果调到 0.3 还是漏检那大概率是训练数据里某个姿态/场景覆盖不足去补一些相似图片加入训练集。误检多往往是背景噪声问题。小黄人以外的物体被当成目标通常是因为训练样本里目标外观太单一模型把“黄色的圆形区域”当成了小黄人。这种时候可以加入一些不含小黄人的背景图作为负样本让模型知道“没有目标也要输出空结果”。虽然官方训练代码没有默认给你负样本接口但你可以把标签清空、只给 boxes 为空数组的样本加进去多训几轮会明显改善。6.3 第一集结束前的最后一个建议最后想跟你说一个我反复体验过的感受目标检测项目里80% 的时间花在数据和调试上只有 20% 花在模型代码上。你可以把模型结构换成别的检测器但只要能熟练处理数据、看懂 Loss、会排查标注问题就具备了迁移到任何检测任务的能力。如果整个流程跑通之后你想继续深入建议下一步从“评估”开始。你会发现训练 Loss 下降并不完全等于检测效果好因为 mAP 的计算方式、IoU 阈值的设定都会影响你对模型能力的判断。这个我们留到第二篇再讲。先把这个小黄人检测模型玩起来多喂几张图试试调整一下置信度阈值感受一下目标检测的输出到底是什么样比什么都重要。