拓冰建站拓冰建站
首页 / 资讯中心 / 正文

Faster-RCNN目标检测代码全解析:从RPN到RoI Head的工程实现

简介Faster R-CNN 目标检测源码资源包是一份面向计算机视觉学习者和算法复现者的完整工程针对目标检测中的候选区域生成、特征提取与分类回归等核心问题提供了基于 PyTorch 的可运行实现也适合作为毕业论文或科研项目的参考基线。压缩包共含 42 个文件以 31 个 Python 脚本为主体另含备份文件、环境/依赖配置文件和 README 说明文档整体体积仅 137KB。已有 134 人学习浏览说明其在目标检测方向具有一定参考价值。资源内部不仅实现了 Faster R-CNN 主框架、RPN 区域提议网络、ROI Pooling 与分类回归层还集成 ResNet50、MobileNetV2、VGG 等多种骨干网络和 FPN 特征金字塔并附带训练、多 GPU 训练、验证、预测、mAP 曲线绘制、COCO 评估及数据划分等工具脚本便于从环境搭建到自定义数据集训练全流程实践。对于想深入理解目标检测原理或快速启动相关实验的开发者这份小型源码包能节省大量整理时间。1. Faster-RCNN目标检测代码为什么两阶段模型仍然值得细读Faster-RCNN目标检测代码 是目标检测领域最值得拆开读的一类工程实现。相比 YOLO 把候选框回归和类别判断合在同一个前向里Faster R-CNN 把“在哪找目标”和“目标是什么”拆成两个阶段结构上更接近模块化设计backbone 负责提特征RPN 负责给候选框RoI Head 负责精分类与坐标回归。也正因为这样小目标漏检、长尾类别召回低、换主干网络这些需求在代码里都有明确的落点改起来比一体化模型更容易定位。下文按日常调试顺序来写先建立模型代码坐标系再跑通推理接着进训练循环最后落到验证与排错。适合刚跑熟 YOLO、想对比两阶段差异的读者也适合已经在用 Faster R-CNN但还没把 anchor、RoI align 这些概念对应到具体代码行的工程师。文中涉及的可执行示例均以 torchvision 的实现为基础这也是目前最主流的二次开发起点。2. 看懂Faster-RCNN代码的模型骨架backbone、RPN与RoI Head2.1 用模型打印建立全图坐标很多人开始改代码时会对着论文里的结构图去读代码效率不高。torchvision 版的 Faster R-CNN 虽然是一个 nn.Module但 forward 里同时处理了训练和推理两条路径看打印结果比看结构图更直接。我一般先执行三行代码import torchvision from torchvision.models.detection import fasterrcnn_resnet50_fpn model fasterrcnn_resnet50_fpn(weightsNone, num_classes91) model.eval() for name, module in model.named_children(): print(name, type(module).__name__)输出从上到下是backbone、rpn、roi_heads、transform前三个是模块第四个负责输入预处理和坐标还原。BackboneWithFPN输出多尺度特征字典在后续代码里统一叫featureskey 是字符串0到4rpn吃掉特征字典后给出一批proposalsroi_heads再根据 proposal 从特征字典中抠出对应区域做分类和回归。理解这份代码不能按nn.Sequential的层序去读要按数据流读。forward的入口是transform它做归一化、resize 并记录缩放比例训练分支还会在这里把 GT 的 boxes 坐标同步缩放。代码里有一个非常隐蔽的分支训练时传入model(images, targets)才有 loss 返回不传targets会静默进入推理分支。如果写训练循环时把 targets 拼错 keyloss 不会报错但一直不更新这类问题靠看日志里的 loss 数值变化才能发现。2.2 锚点生成与分配RPN 代码里最容易被改坏的参数锚点不需要手动提供代码里自动生成这是 Faster R-CNN 和传统 proposal 方法最大的差异之一。下面几行能快速看到默认配置rpn model.rpn print(rpn.anchor_generator.sizes) print(rpn.anchor_generator.aspect_ratios) for name, mod in rpn.head.named_children(): print(name, mod)输出分别是五层锚点尺寸(32, 64, 128, 256, 512)和宽高比(0.5, 1.0, 2.0)head 里有rpn_conv、objectness、pred_bbox_deltas三个卷积。objectness的卷积核数量是 3对应每个位置的 3 种宽高比pred_bbox_deltas的卷积核数量是 12含义是3 种锚点 × 4 个回归量。这里最容易陷入的误判是只有 600 像素的输入分辨率下仍保留 512 这个大锚点它已经超出有效感受野训练阶段 RPN 会被迫匹配这个大尺度目标loss 不降或小目标漏检经常源于此。先确认输入min_size再决定要不要改 sizes 参数。RPN 阶段的正负样本不是直接算 IoU 就完事而是走 Matcher匹配规则默认阈值训练中的作用锚点与 GT 的 IoU 大于 0.7正样本参与 objectness 前景二分类IoU 小于 0.3负样本参与背景二分类IoU 介于 0.3 与 0.7ignore不计 loss不参与参数更新allow_low_quality_matchesTrue保留每组最高 IoU 锚点防止某个 GT 完全没有正样本锚点调参时优先看rpn_positive_fraction默认 0.5。如果训练日志里正样本比例一直偏低问题多半在锚点生成尺度而不是 IoU 阈值上。2.3 RoI Head 的代码路径从 proposals 到分类回归RoI Head 这一段的代码密度最高我同样建议在改代码前先敲一遍roi model.roi_heads print(roi.box_roi_pool) # MultiScaleRoIAlign(output_size(7, 7), sampling_ratio2) print(roi.box_head) # TwoMLPHead(in_channels12544, hidden_layer1024) print(roi.box_predictor) # FastRCNNPredictor(in_channels1024, num_classes91)box_roi_pool就是常说的 RoIAlignoutput_size7表示任意尺寸的 proposal 都会采样成 7×7 特征。sampling_ratio2是每个 bin 的采样点数改成 1 能省一点耗时但小目标回归精度会掉。TwoMLPHead的in_channels是256×7×7256 来自 FPN 的通道数所以换 backbone 时只要输出的通道数不是 256这一层就需要一起改。两阶段在这里体现得最明显第一阶段用自己的定位 loss 训练 RPN第二阶段在裁剪后的特征上重新做 IoU 匹配默认阈值是 0.5和 RPN 阶段的 0.7 是两套监督。第一次读代码的人经常把这两个阈值当成同一个改 RPN 阶段阈值后第二阶段结果完全没有变化原因就在这里。3. 用Faster-RCNN目标检测代码跑推理权重加载与输出解析3.1 最小推理代码和输出格式Faster-RCNN 的推理输出不是 tensor而是 dict坐标也是原图坐标不是归一化结果。下面这段是能在 10 行内跑通的最小示例from PIL import Image import torch from torchvision.models.detection import ( fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights, ) from torchvision import transforms as T weights FasterRCNN_ResNet50_FPN_Weights.COCO_V1 model fasterrcnn_resnet50_fpn(weightsweights) model.eval() device cuda if torch.cuda.is_available() else cpu model.to(device) x T.ToTensor()(Image.open(street.jpg).convert(RGB)).unsqueeze(0).to(device) with torch.no_grad(): pred model(x)[0] print(pred.keys()) # dict_keys([boxes, labels, scores])boxes形状是(N, 4)四个值为[x1, y1, x2, y2]坐标系是原始图片像素位置不是 0 到 1 的归一化区间。这是因为GeneralizedRCNNTransform在postprocess阶段把预测坐标乘回了缩放比例。如果你自己提前把图片 resize 过再喂进网络后续所有面积过滤和可视化都必须回到原图尺寸做。labels是 int64 类别索引不是 one-hot 向量scores是对应类别的置信度。细节上COCO 预训练模型把背景当作 label 0 保留但推理输出里不会出现 0因为背景结果会被丢掉。业务数据如果类别编号从 0 开始转到 COCO 评测格式时整体要加 1。3.2 三个后处理参数score_thresh、nms_thresh、detections_per_img默认后处理参数是为测评准备的不是为线上部署准备的。创建模型时可以这样重写model fasterrcnn_resnet50_fpn( weightsFasterRCNN_ResNet50_FPN_Weights.COCO_V1, min_size640, max_size1200, box_score_thresh0.35, box_nms_thresh0.5, box_detections_per_img300, )参数torchvision 默认值作用线上常用调整box_score_thresh0.05得分低于此值的框被直接丢弃调到 0.4~0.5假正例明显变少box_nms_thresh0.5同类高度重叠框做 NMS 的 IoU 阈值密集堆叠场景降到 0.3box_detections_per_img300每张图最多返回的检测框数小目标密集任务调到 500min_size/max_size800 / 1333输入短边和长边的缩放目标按业务分辨率适当缩小以提速min_size不是目标最小尺寸是图像短边的缩放目标max_size是长边的上限。这两个值对推理耗时的影响比模型结构更直接把max_size从 1333 降到 1024耗时常能降到六成AP 通常只掉零点几个点。box_score_thresh默认很低只有 0.05所以 COCO 权重刚加载直接跑会输出大量低分框看起来模型不行实际只是阈值没提上来。3.3 把模型输出解析成业务协议业务对接时最常用的是 COCO 风格[x, y, width, height]。转换代码建议这样写def pred_to_coco(pred, img_w, img_h): out [] boxes pred[boxes].detach().cpu().numpy() labels pred[labels].detach().cpu().numpy() scores pred[scores].detach().cpu().numpy() for box, label, score in zip(boxes, labels, scores): x1, y1, x2, y2 box x1 min(max(x1, 0), img_w) y1 min(max(y1, 0), img_h) x2 min(max(x2, 0), img_w) y2 min(max(y2, 0), img_h) if x2 x1 or y2 y1: continue out.append({ category_id: int(label), score: round(float(score), 4), bbox: [round(x1, 1), round(y1, 1), round(x2 - x1, 1), round(y2 - y1, 1)], }) return outFPN 输出的框偶尔会在图像边缘产生负坐标或超出宽高的坐标这段代码做了边界截断。另一个容易犯的错是把x2 - x1写反成x1这类问题在评测时不会立刻报错但 mAP 会低得莫名其妙。4. 训练Faster-RCNN代码数据集封装、损失与超参数4.1 自定义数据集和 collate_fn 的写法Faster R-CNN 的 DataLoader 和普通分类任务不太一样它需要的数据结构是List[dict]而不是堆好的张量。每个 target 里必须有boxes和labels两个键from torch.utils.data import Dataset class DetectionDataset(Dataset): def __init__(self, records): self.records records def __len__(self): return len(self.records) def __getitem__(self, idx): rec self.records[idx] # rec 含字段: image(BGR ndarray), boxes(list of XYXY), labels image rec[image][:, :, ::-1].copy() # BGR - RGB boxes torch.as_tensor(rec[boxes], dtypetorch.float32) labels torch.as_tensor(rec[labels], dtypetorch.int64) return image, {boxes: boxes, labels: labels} def collate_fn(batch): return [item[0] for item in batch], [item[1] for item in batch]boxes必须是 XYXY 格式因为内部在算 IoU 时按这个假设写。labels从 1 开始0 被 torchvision 保留给背景如果源数据集的类别从 0 编号要整体加 1。collate_fn返回的是两个 list 而不是 torch.stack 的结果。原因是 batch 内图片尺寸不一定相同torchvision 的GeneralizedRCNNTransform会在 forward 内部统一缩放提前 stack 会直接报形状不一致。反过来如果数据已经裁成固定尺寸用 stack 也能跑但一旦之后加入随机裁剪就会现原形所以统一用宽松写法最省事。4.2 训练主循环返回的是 loss 字典一个能正常收敛的训练循环骨架如下from torch.optim import SGD from torch.optim.lr_scheduler import StepLR params [p for p in model.parameters() if p.requires_grad] optimizer SGD(params, lr0.005, momentum0.9, weight_decay0.0005) lr_scheduler StepLR(optimizer, step_size3, gamma0.1) for epoch in range(epochs): model.train() for images, targets in dataloader: images [image.to(device) for image 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() lr_scheduler.step() print({k: round(v.item(), 3) for k, v in loss_dict.items()})model.train()和model.eval()在这里直接影响 RPN 采样的 proposal 数量不能省。model(images, targets)返回的字典有四个 keyloss 名称作用数值异常时优先检查loss_objectnessRPN 的前景/背景二分类锚点尺度和输入分辨率不匹配loss_rpn_box_regRPN 的锚点坐标回归目标框坐标存在 NaN 或没有归一化loss_classifierRoI Head 的类别预测类别索引从 0 开始被当成背景loss_box_regRoI Head 的坐标精修labels 值超过了num_classes训练时要把这四项单独打出来只看 total loss 会掩盖内部问题。比如loss_objectness已经稳定下降但loss_rpn_box_reg始终在抖说明前景框找到了但坐标没学稳这时候去调rpn_positive_fraction比盲目降学习率更有效。4.3 训练必调的 4 个超参数与梯度累加超参数起调值影响常见误用lr0.005SGD整体收敛速度大于 0.02 时 RPN loss 剧烈震荡batch_size2显存占用和 BN 统计batch1 也能跑但每个 step 差异大box_positive_fraction0.25RoI Head 采样中正样本比例单类别任务可以设到 0.5step_size3学习率下降周期数据量少于几千张时 step3 偏快显存不够但想用更大 batch 时用梯度累加。需要注意累加逻辑是把多个 batch 的梯度累积后做一次更新学习率仍按大 batch 去理解accum_steps 4 model.zero_grad() for step, (images, targets) in enumerate(dataloader): 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) loss sum(v for v in loss_dict.values()) / accum_steps loss.backward() if (step 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()除以accum_steps是为了让累计梯度量级和真正的大 batch 一致。这里不能偷懒只对一个 head 的 loss 累加四个单项的 scale 各不相同单独加权反而容易让某个 loss 占主导。5. 验证与排错把Faster-RCNN目标检测代码调到可用状态5.1 训练间隙做快速数值检查不用等完整 mAP训练时每轮看几个数值就能判断模型是否进入有效学习。代码片段如下model.eval() with torch.no_grad(): pred model(images)[0] keep pred[scores] 0.5 num_dets int(keep.sum()) boxes_keep pred[boxes][keep] w boxes_keep[:, 2] - boxes_keep[:, 0] h boxes_keep[:, 3] - boxes_keep[:, 1] ratios w / (h 1e-6) print(num_dets, float(ratios.mean()), float(pred[scores][keep].mean()))num_dets很低说明阈值太高或模型在验证图上基本无响应ratios明显偏离业务目标的宽高比说明回归还没收敛mean score长期接近 1则要警惕过拟合。把这组检查放到验证集前 20 张图上一轮只要几秒。5.2 三个高频问题的定位方向现象定位处理训练时 loss 下降正常验证集框全空验证时没调用model.eval()或用了训练模式下的 RPN 采样验证入口先切换模型状态换自定义数据集后 loss 为 NaNboxes 里有空框或负宽高坐标 dtype 不统一在 Dataset 的__getitem__里过滤掉x2x1的框单类别数据 mAP 特别低类别索引从 0 开始被当成背景统一把 labels 加 1输出时再减回来5.3 ONNX 导出时最容易踩的坑Faster R-CNN 推理输出的boxes数量每张图都不一样直接用torch.jit.trace导出时容易把 batch 维度或候选框数量记成固定值换一张图就报 shape 不匹配。常见做法是先把输入固定成(1, 3, 800, 1200)再 trace并在导出前把box_score_thresh调高减少输出框数量的浮动范围。保存权重之后第一件事是用同一张训练图跑一遍原始 PyTorch 模型和导出模型对比pred[boxes][:5]是否一致如果坐标对不上问题通常在预处理里的均值方差没有固化进去。部署时不要在每一次请求里重建模型把model.eval()和权重加载放到进程启动时模型会一直复用同一份显存输出坐标也稳定一致。本文还有配套的精品资源点击获取
分享:

看完干货,该让你的企业上线了

免费需求沟通 · 48 小时内出具建站方案 · 河南本地可上门