Faster R-CNN之rcnn_head全解析:从RoIAlign到分类回归 写这个faster rcnn代码解读系列写到第九篇终于轮到rcnn_head了。前八篇我们把backbone、FPN、RPN这些外围模块都拆完了整个检测框架其实已经把可能有物体的proposal“找”了出来但这些proposal到底对应什么类别、框的位置要不要再微调一下都还悬着。rcnn_head就是负责解决这个问题的模块它把RPN给到的proposal和backbone加FPN产出的多尺度特征图拉到一起通过RoIAlign抠出每个候选框的区域特征再经过两层全连接做分类和框回归最终输出“这个框里是什么、该怎么修正”。一句话概括rcnn_head是整条检测流水线上真正出结论的地方。如果你是第一次读这套代码推荐先把前八篇快速过一遍尤其是FPN那篇因为rcnn_head的RoIAlign部分和FPN特征层选择是绑在一起的如果你只想彻底搞懂head本身这篇可以直接当作独立章节看。我会从整体流程、关键类结构、loss计算、推理细节、常见坑这几个方向展开尽量把代码背后的设计理由也讲清楚。1. 先把rcnn_head放进整条检测链路里看1.1 前向过程中rcnn_head的输入与输出很多初学者看rcnn_head时容易一上来就钻进Shared2FCBBoxHead的代码里结果被一堆Linear和ModuleList绕晕。我的建议是先搞清楚这个模块在整个前向流程中的位置再回头看实现细节。训练或推理时整张图片先经过backbone得到C2到C5的特征再经FPN融合成P2到P5的多尺度特征图。RPN在这些特征图上做区域提议输出一批proposal格式是原图坐标系下的xyxy坐标同时附一个objectness分数。到这里模型只是知道“这里可能有个东西”但不知道“这是猫还是狗”。rcnn_head就是接在RPN后面的模块它接收两个东西featuresFPN输出的多尺度特征图通常是一个dict比如{p2: (B,256,H/4,W/4), p3: (B,256,H/8,W/8), p4: (B,256,H/16,W/16), p5: (B,256,H/32,W/32)}。proposalsRPN输出的候选框通常有几千个推理时会按score排序后截断比如只取前1000个。它的输出分两种情况。训练时输出两个lossloss_cls和loss_bbox用来更新整个模型包括backbone和RPN因为Faster R-CNN是端到端训练的。推理时输出最终的检测结果也就是每个框的类别、置信度和修正后的坐标。这里有个容易混淆的点RPN的输出也叫“proposal”或“roi”rcnn_head内部又把它们叫“rois”。其实是一个东西。代码里经常出现num_proposals这个参数它控制每次送入head的候选框数量训练时默认512测试时可能是1000。这个数量很关键数量太少会漏检数量太多显存吃不消后面会展开讲。1.2 从代码看RCNNHead类的内部结构我在这个系列里统一读的是基于mmdetection 2.x风格的代码rcnn_head的核心类叫RCNNHead文件路径一般是models/roi_heads/roi_head.py。类的定义看起来很长但骨架并不复杂核心逻辑可以缩成这样class RCNNHead(nn.Module): def __init__(self, num_stages1, num_proposals512, with_bboxTrue): super().__init__() self.num_stages num_stages self.num_proposals num_proposals self.with_bbox with_bbox if with_bbox: self.bbox_roi_extractor SingleRoIExtractor( roi_layerdict(typeRoIAlign, out_size7, sampling_ratio2), out_channels256, featmap_strides[4, 8, 16, 32]) self.bbox_head Shared2FCBBoxHead( in_channels256, fc_out_channels1024, roi_feat_size7, num_classes80) def forward(self, features, proposals, gt_bboxesNone, gt_labelsNone): if self.training: return self.forward_train(features, proposals, gt_bboxes, gt_labels) return self.forward_test(features, proposals)RCNNHead里面并不直接写分类和回归的计算而是把工作拆给两个子模块bbox_roi_extractor负责从多尺度特征图上抠出proposal区域的特征类型是SingleRoIExtractor内部配置了RoIAlign池化层。它的输出是一个固定形状的tensor方便后面全连接层处理。bbox_head负责真正的分类和回归类型是Shared2FCBBoxHead。名字里的“Shared2FC”指两个全连接层是共享的也就是cls分支和reg分支先共用两层FC然后再各自接一个输出层。为什么要拆成两个子模块因为职责不同。第一个模块解决“怎么把任意形状的proposal变成统一形状的特征”第二个模块解决“从特征中预测类别和坐标偏移”。拆开之后如果想换一个池化方式或换一个head结构只需要替换对应子模块其他代码不用动。这种面向接口的设计在检测框架里非常常见读代码时抓住这个分层思路会轻松很多。2. bbox_roi_extractor怎么把proposal变成统一尺寸的特征2.1 RoIAlign与FPN的配合很多人对RoIAlign的印象停留在“双线性插值解决量化误差”但在FPN场景下还有一个更实际的问题proposal大小不一该从哪一层特征图上抠特征这就是bbox_roi_extractor的核心逻辑。先看配置bbox_roi_extractordict( typeSingleRoIExtractor, roi_layerdict(typeRoIAlign, out_size7, sampling_ratio2), out_channels256, featmap_strides[4, 8, 16, 32])这里的featmap_strides对应FPN各层的下采样倍数。proposal的坐标是原图尺度要映射到特征图就得除以对应层的stride。比如一个200x300的框在stride16的层上就是12.5x18.75在stride4的层上就是50x75。不同尺度的目标适合在不同分辨率的层上检测小目标在浅层高分辨率特征上信息保留更完整大目标在深层低分辨率特征上语义更清晰。那怎么决定一个proposal用哪一层代码里通常按面积来选。规则大致是先算proposal的宽高然后根据宽高乘积映射到某个层。比如宽高都小于某个阈值就用stride较小的层大于某个阈值就用stride较大的层。MultiScaleRoIAlign里还有一个finest_scale参数默认56意思是如果proposal的短边小于56就认为它属于最精细的特征层。RoIAlign本身做的事可以理解成“把特征图上任意一个矩形区域切出来再缩放到固定大小”。假设我们想输出7x7的特征块就把roi区域均匀切成7x7的网格每个格子称为一个bin再在每个bin内部做双线性插值采样。sampling_ratio2表示每个bin里采样2x2个点然后取平均或最大值作为这个bin的输出。为什么要花这么多力气做浮点对齐因为RoIPooling有两次量化操作一次是把浮点坐标取整到整数网格一次是把roi区域均分时再次取整。对小目标来说一次取整就可能丢掉几个像素累积误差很致命。RoIAlign把坐标和采样点都保留浮点精度用双线性插值取特征值能明显提升小目标检测框的准确性。这也是rcnn_head最终输出质量的基础保障。2.2 一个7x7特征块是怎么来的我们用一个具体例子把数据流走一遍这样印象更深。假设输入图片是800x1000FPN输出了p2到p5四层特征stride分别是4、8、16、32。RPN给了一个proposal原图坐标是[100, 150, 400, 450]宽300、高300面积9万像素。按面积或宽高判断这个proposal大概率落在stride16的特征层上。坐标除以16得到[6.25, 9.375, 25.0, 28.125]。RoIAlign把这个浮点矩形区域分成7x7个bin每个bin内部采样2x2个点从对应特征图上取双线性插值结果。最终这个proposal得到一个形状为(256, 7, 7)的特征块256是FPN输出的通道数。一批处理时假设有512个proposal那么输出就是(512, 256, 7, 7)的tensor。后面要做的就是把每个7x7的特征块展平成256*4912544维的向量再送进全连接层。这个地方有个很常见的疑问为什么out_size选7而不是14或28out_size越大每个bin的采样点越多空间信息保留得更细但全连接层的输入维度会急剧膨胀。7x7是Fast R-CNN时代流传下来的经典配置在精度和计算量之间取得了平衡。我试过改成14x14精度确实能涨一点但显存占用和训练时间也明显增加对于大多数场景并不划算。3. Shared2FCBBoxHead分类回归分支的构建细节3.1 两层shared FC到底在做什么拿到(512, 256, 7, 7)的特征后接下来进入bbox_head。这个head的名字叫Shared2FCBBoxHead核心结构是两层共享全连接加两个输出分支。代码可以简化为self.shared_fcs nn.ModuleList() # 第一层把展平特征映射到1024维 self.shared_fcs.append(nn.Linear(256 * 7 * 7, 1024)) # 第二层再做一次1024到1024的映射 self.shared_fcs.append(nn.Linear(1024, 1024)) self.fcs_cls nn.ModuleList() self.fcs_reg nn.ModuleList() # 分类分支 self.fcs_cls.append(nn.Linear(1024, num_classes)) # 回归分支 self.fcs_reg.append(nn.Linear(1024, num_classes * 4))其中fc6、fc7这两个叫法是从VGG16时代传下来的意思是第6、第7个全连接层。虽然ResNet替代了VGG但代码里还是常用fc6、fc7来指代这两层读别人代码时会经常看到这种命名。为什么分类和回归要共享前两层全连接一个重要原因是分类和回归本质上都在分析同一个区域的特征低层信息是通用的。共享特征能减少参数量也能让两个任务互相补充。但如果完全共享也容易导致两个任务的梯度互相干扰所以最后必须分成两个独立输出层让每个任务有自己的“出口”。有一点需要注意Shared2FCBBoxHead里其实还预留了shared_convs的支持也就是先用卷积层做特征提取再接全连接。但ResNet50加FPN的配置通常不用卷积直接共享两层FC就够了。对很多中小型目标检测任务来说两层FC的表达能力已经足够再加深边际收益很低反而拖慢训练和推理速度。3.2 cls_score与bbox_pred的维度设计在理解维度设计前先明确一个概念分类分支和回归分支不是平等的输出关系。分类要回答“这个框里是什么”输出的是每个类别的置信度回归要回答“这个框怎么修”输出的是proposal与真实框之间的坐标偏移。看维度cls_score形状是(N, num_classes)N是送入head的proposal数量。有的实现会多一个背景类有的不会。当前这个版本在数据加载时已经把背景类处理掉了所以num_classes就是真实类别数比如COCO的80类。bbox_pred形状是(N, num_classes * 4)。为什么每个类别都要单独预测一组偏移因为不同类别的长宽比、整体大小差异很大。比如“人”一般是瘦长的“车”一般是扁宽的如果所有类别共用一组回归参数模型就得在类别间妥协精度会受影响。为每个类别维护一组偏移相当于让模型“知道”我要修正的目标是什么形状。训练时回归loss只针对正样本的对应类别来算。假设某个proposal被分配给了第5类就只取bbox_pred里第5类对应的那4个值跟该proposal的回归target做smooth L1计算。负样本不参与回归loss因为它的“目标类别”不存在强行回归没有意义。3.3 label分配的细节head在训练时不能随便拿512个proposal去算loss得先告诉模型“这里面哪些是正样本、哪些是负样本、每个正样本对应哪个类别、回归目标是什么”。这个环节在代码里叫assign_targets是整个rcnn_head训练逻辑里最容易出bug的地方。通常做法是从RPN给出的全部proposal中随机采样一部分比如512个。把每个proposal和所有gt框做IoU计算。如果某个proposal和某个gt框的IoU大于0.5就标记为正样本label设为该gt框的类别回归目标设为该proposal相对于gt框的偏移。如果IoU在0.1到0.5之间标记为负样本label设为背景通常是0或-1。IoU低于0.1的proposal直接丢弃不参与训练因为它们太难属于无效负样本。这里有个设计取舍值得说为什么阈值是0.5而不是更高0.5意味着proposal和gt框的重叠程度超过一半就认为它“足够像”这个物体。阈值太高比如0.7正样本数量会骤减因为RPN给出的框本来就不算特别准阈值太低比如0.3会把大量半遮半掩的背景框当成正样本引入噪声。0.5是经过大量实验验证的平衡点。正负样本比例也很重要。代码里用pos_balance0.25来控制意思是希望采样出的512个proposal里正样本约占25%也就是128个左右。假设正样本只匹配到80个那负样本也不会强行拉满到432个而是会适当地少采一些保证总样本数可控。我实际调试时遇到过正样本特别少的情况比如某些类别的目标只有几个最后512个proposal里正样本只有二三十个训练明显不稳定。这种情况下我会调整阈值或增加每张图的proposal采样数而不是硬着头皮训。注意这里的0.5是最终的分类判定阈值和RPN训练时的前景阈值概念不同。RPN的阈值是为了筛选“可能有物体的区域”rcnn_head的阈值是为了确定“这个区域到底是什么物体”不要混在一起。4. 损失函数cls loss和reg loss怎么对齐4.1 两个损失的计算过程rcnn_head的loss计算可以缩成两行loss_cls F.cross_entropy(cls_score, labels) loss_bbox smooth_l1_loss(bbox_pred, bbox_targets, beta1.0)但实际代码里还藏了一些细节比如回归loss要考虑正负样本的mask。一个常规的写法是cls_score, bbox_pred self.bbox_head(roi_feats) loss_cls F.cross_entropy(cls_score, labels) # 只取正样本的回归结果 pos_inds (labels 0).nonzero().squeeze(1) bbox_pred_pos bbox_pred[pos_inds] bbox_targets_pos bbox_targets[pos_inds] loss_bbox smooth_l1_loss(bbox_pred_pos, bbox_targets_pos, beta1.0, reductionmean)labels是所有采样proposal的类别编号正样本是1到80负样本是0。cross_entropy会对所有512个proposal算loss包括负样本因为模型需要学会把负样本判断为背景。这就是分类任务对负样本的处理逻辑不是忽略而是告诉模型“这是背景”。回归loss则不同。负样本的回归target是无效的所以必须过滤掉。这里有个容易踩的坑如果你在构造bbox_targets时没把负样本的回归target置零或没做masksmooth L1会把垃圾target也一起算进去loss会变得很奇怪。4.2 为什么用它 / 怎么调loss_weight先看smooth L1的公式当|x| 1时loss 0.5 * x^2当|x| 1时loss |x| - 0.5这个函数的设计逻辑很巧妙。如果用纯L2损失碰到离群点也就是预测偏差特别大的样本时梯度会非常大一更新参数就震荡训练容易发散。如果用纯L1损失梯度恒定为1离群点不会引爆训练但在0附近梯度不连续精度会受限。smooth L1把两者结合偏差小的时候走L2的平滑曲线偏差大的时候走L1的稳定梯度既抗离群点又保证局部精度。那loss_weight怎么配默认都是1.0也就是分类和回归对最终loss的贡献各占一半。我实际经验是如果发现模型“类别经常分错但框还算准”可以适当加大loss_cls的权重比如1.5比1如果模型“框总是飘但类别挺准”就加大loss_bbox的权重。但调整幅度别太大1比1到2比1之间足够覆盖大部分场景再往上调很容易顾此失彼。这里再分享一个实用技巧训练时别只看总loss要把loss_cls和loss_bbox分别打出来看走势。我遇到过一种情况总loss一直在降但loss_bbox反而在涨说明分类任务在主导优化方向回归任务正在被带偏。这种问题如果不单独看分项loss很难发现。5. 推理阶段从score阈值到最终nms5.1 推理时怎么去掉不靠谱的框推理时RCNNHead走forward_test分支逻辑比训练简单但细节更多。首先它会把RPN给出的所有proposal按score排序只保留前num_proposals个比如1000个。然后把所有proposal一次性丢给bbox_roi_extractor和bbox_head拿到cls_score和bbox_pred。接下来是后处理流程对cls_score做softmax或sigmoid取决于训练时用的是交叉熵还是focal loss得到每个框属于每个类别的概率。设定score_thr默认0.05把得分低于这个值的框直接丢掉。这一步能滤掉大量低置信度的背景框。对保留下来的框用bbox_pred里对应类别的4个值做解码也就是把RPN给出的原坐标加上预测偏移得到修正后的坐标。按类别做nmsIoU阈值默认0.5确保同一个物体不会被多个高度重叠的框框住。最后按得分排序只保留前max_per_img个结果输出。一个很常见的调参场景是如果最终结果里有很多框重叠在一起先别急着怀疑模型没训练好先看两件事。一是score_thr是不是设得太低导致大量低分框进入了nms阶段二是nms的IoU阈值是不是太宽松比如设成0.7重叠框就很难被抑制掉。我一般会把score_thr设在0.05到0.1之间nms阈值设在0.5附近。5.2 torchvision nms替换旧实现想特别说一下nms的版本变化。以前mmdetection自带一个叫lib_nms的C/CUDA实现性能也不错但问题是它要参与项目编译环境一变就各种编译报错非常折腾。从0.4.0版本开始代码里换成了torchvision自带的batched_nms直接开箱即用版本升级也方便。batched_nms支持一次性对多个类别做nms。它的做法是给每个框类别加上一个很大的偏移量把不同类别“拉开”然后统一排序和抑制。这样就不用写一个for循环重复调nms效率更高。代码里会看到一个class_agnostic参数。class_agnosticTrue表示所有类别放在一起做nmsFalse表示按类别分开做。正常推理一般用False因为不同类别的目标即使高度重叠也应该保留比如一个人骑着一辆车。True的语义更像“不管是什么类别只要框重叠就干掉”通常用于hard negative mining这类特殊场景。提示如果发现更新torchvision后nms结果和旧版不完全一致不用太纠结。batched_nms在跨类别时不保证每类单独nms后合并结果完全相等但对最终检测精度的影响通常很小。5.3 推理结果的可视化验证代码写完之后最重要的一步是可视化验证。拿一张测试图跑一遍推理把最终结果画在图上先看类别对不对、框紧不紧。如果类别正确但框明显偏大或偏小优先调整nms阈值和回归loss的权重如果同一个物体出现大量重叠框优先调低score_thr或者调严nms阈值。很多人在这一步容易犯的错是只看mAP数字不直接看图。mAP高不代表每个类别都好尤其在小目标或密集场景下可视化能暴露很多指标看不出的问题。我把proposal可视化之后经常发现模型在密集人群里把两个紧挨着的人框成了一个这种问题靠调score_thr和nms阈值不一定能完全解决得从锚点生成和nms策略入手。6. 实操复盘用hook把中间过程“看”出来6.1 注册forward hook取中间特征rcnn_head是个黑盒子直接看loss曲线很难判断中间特征长什么样。我强烈建议在debug阶段给bbox_head注册一个forward hook把中间层的输出打出来看看。简单示例def hook_fn(module, input, output): cls_score, bbox_pred output print(cls_score shape:, cls_score.shape) print(bbox_pred shape:, bbox_pred.shape) print(cls_score max/min:, cls_score.max().item(), cls_score.min().item()) bbox_head model.roi_head.bbox_head bbox_head.register_forward_hook(hook_fn) with torch.no_grad(): model(batch)hook能打出当前模型对这批数据的中间判断比如所有proposal的类别分数分布、回归偏移的大致范围。我以前习惯在训练时每跑几百个iter打印一次看看score_min和score_max的变化趋势。如果score_max迟迟不接近1说明模型对“最像正样本的proposal”都没太大把握大概率是正样本太少或特征没有对齐。6.2 拿到数据后怎么看拿到中间输出后除了看shape更重要的是做一些针对性统计。统计cls_score里最大值的分布如果大部分proposal的最大类别分数都在0.1以下说明模型还没学会“信任”任何proposal训练可能还在早期如果训练了很久还是这样就要怀疑正负样本比例或学习率。统计bbox_pred的范围回归偏移过大比如超过10说明回归target没做归一化或者proposal和gt差异太大。合理情况下偏移量应该在0附近小幅波动因为RPN给出的proposal已经和gt比较接近了。还可以把hook里拿到的proposal坐标直接画原图上看看哪些区域的proposal被大量保留、哪些区域被丢弃。这能直观地暴露漏检区域比如某些被遮挡物体根本没被proposal覆盖到。我举个例子有次调一个密集小目标场景的模型loss曲线一直很漂亮但测试时漏检严重。挂上hook一看发现大量小目标的proposal被score_thr过滤掉了因为它们几乎都低于阈值。问题根本不是分类头没训练好而是proposal本身的特征提取在小目标上太弱后来换了更大的out_size并调整了FPN的层匹配情况才好转。7. 常见问题与排查实录现象可能原因排查方法分类全错但框比较准label分配对不上类别映射错位打印labels分布检查gt标注类别范围和模型输出类别是否一致loss不降或下降极慢正负样本严重不均衡或学习率不合适检查pos_balance和采样数打印正负样本数量尝试调整基础学习率最终结果大量重叠框score_thr太低或nms阈值太大降低score_thr到0.05左右nms阈值调到0.5附近bbox预测出现nan或inf回归target没限制gt框或proposal坐标越界检查bbox_targets的构造对坐标做clip或归一化torchvision编译或nms报错版本太旧缺少batched_nms升级torchvision到对应版本确认nms调用路径正样本数量总是个位数gt标注漏检或proposal覆盖质量太差可视化RPN输出看proposal是否覆盖到目标区域train loss正常但test效果差过拟合或head的共享特征表达能力不足检查训练集和测试集数据分布考虑换更大的head结构再补充几条避坑心得。第一不要随意调大num_proposals它直接影响显存占用和训练速度。512是一个比较稳的默认值如果你的输入图特别大可以适当降到256但不要为了提高召回硬塞到1024很容易OOM。第二修改roi_feat_size和fc_out_channels后要同步检查bbox_head里全连接层的输入维度比如256*7*712544这个数字是不能随便变的一改整个head都要跟着改。第三推理时用的score_thr和训练时的阈值没有直接关系训练时0.5只是用来分配正负样本推理时0.05是最终输出过滤两个概念别搞混。我个人在这一块踩过不少坑最值得说的一个体会是rcnn_head是整个检测器里最容易出现“看着loss正常、实际预测不对”的地方。因为分类和回归共享特征两个任务的梯度会互相影响调参的时候一定不能单看总loss要把loss_cls和loss_bbox分开打出来看趋势。还有一个很实用的小技巧训练中期把proposal可视化一次很多标签错误、对齐问题一眼就能看出来比盯着曲线猜有效得多。这个系列到现在已经把Faster R-CNN的主干、FPN、RPN和head都过了一遍了下一步我准备再写一篇关于单阶段检测器里head的差异正好和这篇形成对比把两套方案的取舍讲透。