YOLO改进：残差注意力模块RAB嵌入C3k2，提升小目标召回率

目标检测的模型改进现在基本都围着主干、颈部、头部三个结构打转而注意力模块是其中最耐用的“膏药”。这个系列走到第77期我这次要记录的是一个自定义的RABResidual Attention Block残差注意力模块把它嵌到C3k2模块内部形成带注意力分支的瓶颈结构。先说结论整个模块增加参数量很小但精度收益比直接堆SE、CBAM更稳尤其是对小目标和遮挡场景召回率提升比较明显。如果你也在做YOLO系模型的轻量结构改造这篇文章可以直接抄作业。我会把设计思路、完整代码、插入C3k2的改法、训练配置、踩坑记录都铺开讲。代码基于PyTorch不需要改YAML主结构只加一个自定义模块类就能跑。整个过程我在自定义的20类数据集上复现过也拆过训练日志目前来看RAB放在瓶颈内部、只作用于深层Stage的收益最高全图无脑插入反而容易震荡。1. 为什么在C3k2里加RAB而不是继续叠SE1.1 从C3k2的结构看改进的切入点C3k2这类模块本质上是“CSP风格的多瓶颈堆叠”输入先过一个1x1卷积中间串若干个Bottleneck侧边再来一条1x1卷积最后拼起来过一层输出卷积。它相比早期C3结构最大的变化是梯度流的分配方式更合理Bottleneck自带的shortcut也缓解了深层信息丢失的问题。但这里有个容易被忽略的盲区Bottleneck的shortcut只能保证“梯度能流回去”它不能建模通道之间的依赖关系。不同特征通道在检测任务中的重要性是不同的比如车辆检测里车的轮廓通道和轮子纹理通道都很关键但它们的激活强度可能差一个量级。传统做法是加SE模块先在通道维度做全局池化再重新标定。SE本身没问题可一旦放在C3k2这种分辨率已经缩小到16x16、8x8的深层Stage时直接乘上0到1之间的通道权重会让整个分支的输出尺度发生变化后续BN统计很容易被带偏。我一开始就是按这个思路试的Loss在头20个epoch里掉得很快但到了第120个epoch就开始抖动。RAB的设计目标很明确注意力分支负责筛选重要特征残差连接负责稳住尺度。注意力输出多少就只对分支局部做调整原始信息始终保持一份恒等映射。这样既拿到了注意力带来的特征筛选能力又不会破坏深层特征的数值分布。1.2 残差注意力跟普通注意力的差别在哪普通注意力比如SE、CBAM、SimAM本质都是“用某种统计量生成0到1的权重再乘到特征图上”。这个操作的问题是权重小于1的时候特征整体被压缩多次堆叠之后特征幅度会被逐层缩小如果某个通道被分配到接近0的权重那这个通道的梯度在反向传播时也会被压得很弱网络很难把它再激活回来。RAB的核心差别是引入了一个恒等旁路。我用的结构是“attention gate residual add”[ out x \cdot M_c \cdot M_s x ]其中 (M_c) 是通道注意力权重(M_s) 是空间注意力权重两者都在0到1之间。输出不再是被压缩后的特征而是“原特征被注意调制的特征”。实验里最直观的表现是RAB插入后训练初期的梯度范数比纯CBAM更稳定因为残差旁路保证至少有一半的梯度可以走恒等路径直接回传。但这里有个实现细节容易翻车如果直接初始化sigmoid输出为0.5那第一轮前向就是 (x*0.5 x 1.5x)等于把特征幅度强行拉高了50%深层网络很快就会爆。所以我在RAB里专门做了初始化让最后一层1x1卷积和空间卷积的bias初始化为-3权重全0这样sigmoid(0)之前输出约0.047初始状态几乎等于恒等映射。这一步非常重要很多复现RAB翻车都是栽在初始数值上。1.3 RAB在C3k2里的最佳位置同一个RAB放的位置不同效果天差地别。我试过三种位置A整个C3k2模块外再接一个RAB。简单粗暴但粒度太粗注意力只能看到模块整体输出感知不到内部哪个Bottleneck提取的特征有价值。位置B每个Bottleneck的第二个卷积输出后接RAB也就是我最终采用的方式。注意力在每层瓶颈内部生效可以逐级筛选。位置CC3k2末尾的concat拼接后、输出卷积之前接RAB。参数增量最小但我实测精度提升只有位置B的一半左右。原因也不难理解。C3k2内部的Bottleneck是串行关系前一个瓶颈的输出是后一个瓶颈的输入。RAB放在每个Bottleneck尾部相当于每一层都在“清洗”上一层的噪声特征逐层累积的效果肯定比只在末尾加一次要强。代价是参数量会随着n增加线性增长所以我的建议是只改造深层Stage的C3k2P3层以前的浅层结构不要动。2. RAB模块设计与代码实现2.1 RAB的完整设计结构RAB不是一个单一的卷积块而是“通道注意力空间注意力残差融合”的组合体。通道注意力部分用全局平均池化提取每个通道的全局描述再用两个1x1卷积做非线性映射最后过sigmoid得到与输入通道数相同的权重向量。我不建议用全连接层做这个映射因为全连接依赖输入特征图的宽高固定而1x1卷积对任意分辨率都友好在测试阶段换输入尺寸时不会报错。空间注意力部分则是取通道方向上的均值和最大值拼成两个通道的“空间描述”再用一个大卷积核卷积得到空间权重。大卷积核我选了7x7因为空间注意力的感受野越大对遮挡目标的容错性越好。小卷积核比如1x1或3x3只能看到像素级局部关系很容易被背景噪声干扰。融合阶段用乘法把通道权重和空间权重逐级作用在原始特征上最后加一条恒等路径。整体计算量很低对一个512通道的Bottleneck来说RAB新增参数量约0.033M其中大头是通道注意力那两个1x1卷积空间卷积只有2进1出的7x7卷积参数量可以忽略不计。2.2 RAB的PyTorch实现下面是我在项目里实际用的完整实现可以直接复制到自定义模块文件里。import torch import torch.nn as nn class ConvBnSiLU(nn.Module): 标准卷积块卷积 BN SiLU兼容常规YOLO代码库的风格 def __init__(self, c_in, c_out, k1, s1, p0): super().__init__() self.conv nn.Conv2d(c_in, c_out, k, s, p, biasFalse) self.bn nn.BatchNorm2d(c_out) self.act nn.SiLU(inplaceTrue) def forward(self, x): return self.act(self.bn(self.conv(x))) class RAB(nn.Module): Residual Attention Block 通道注意力 空间注意力 残差恒等路径 def __init__(self, in_channels, reduction16, spatial_k7): super().__init__() hidden max(8, in_channels // reduction) # 通道注意力AvgPool - 1x1卷积 - ReLU - 1x1卷积 - Sigmoid self.channel_att nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, hidden, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(hidden, in_channels, 1, biasTrue), nn.Sigmoid() ) # 空间注意力拼接平均池化和最大池化 - 7x7卷积 - Sigmoid self.spatial_att nn.Sequential( nn.Conv2d(2, 1, kernel_sizespatial_k, paddingspatial_k // 2, biasTrue), nn.Sigmoid() ) self._init_weights() def _init_weights(self): # 关键初始化让初始注意力权重接近0保证RAB初始接近恒等映射 last_ch_conv self.channel_att[3] nn.init.zeros_(last_ch_conv.weight) nn.init.constant_(last_ch_conv.bias, -3.0) last_sp_conv self.spatial_att[0] nn.init.zeros_(last_sp_conv.weight) nn.init.constant_(last_sp_conv.bias, -3.0) def forward(self, x): identity x # 通道注意力分支 ch_weight self.channel_att(x) # (B, C, 1, 1) x_ch x * ch_weight # 空间注意力分支 avg_map torch.mean(x_ch, dim1, keepdimTrue) max_map, _ torch.max(x_ch, dim1, keepdimTrue) sp_weight self.spatial_att(torch.cat([avg_map, max_map], dim1)) out x_ch * sp_weight return identity out这段代码里有几个细节必须说明白。首先是通道注意力最后一层卷积的bias必须存在因为我们要用负偏置把sigmoid的输出压到接近0从而让初始模块表现为恒等映射。如果这里biasFalse初始权重全0则sigmoid(0)0.5第一次前向就会让特征幅度放大50%。其次是空间注意力最开始的输入是通道注意力之后的特征图不是原始输入这样空间注意力能感知到已经被筛选过的信息两个注意力机制是级联的而不是简单的并行关系。2.3 改造Bottleneck与C3k2RAB本身只是一个组件真正要替换的是C3k2里的Bottleneck。做法是定义一个Bottleneck_RAB在原有“两段1x1/3x3卷积”之后接RAB。import torch.nn as nn class Bottleneck_RAB(nn.Module): 带RAB的Bottleneck卷积提取 - RAB注意力调制 - 残差连接 def __init__(self, c1, c2, shortcutTrue, g1, k(3, 3), e1.0, reduction16): super().__init__() c_ int(c2 * e) self.cv1 ConvBnSiLU(c1, c_, k[0], 1, k[0] // 2) self.cv2 ConvBnSiLU(c_, c2, k[1], 1, k[1] // 2) self.rab RAB(c2, reductionreduction) self.add shortcut and c1 c2 def forward(self, x): y self.cv2(self.cv1(x)) y self.rab(y) return x y if self.add else y class C3k2_RAB(nn.Module): C3k2变体内部所有Bottleneck替换为Bottleneck_RAB def __init__(self, c1, c2, n1, c3kFalse, e0.5, g1, shortcutTrue, reduction16): super().__init__() c_ int(c2 * e) self.cv1 ConvBnSiLU(c1, c_, 1, 1, 0) self.cv2 ConvBnSiLU(c1, c_, 1, 1, 0) self.cv3 ConvBnSiLU(2 * c_, c2, 1, 1, 0) self.m nn.ModuleList( Bottleneck_RAB(c_, c_, shortcut, g, k(3, 3), e1.0, reductionreduction) if c3k else Bottleneck_RAB(c_, c_, shortcut, g, e1.0, reductionreduction) for _ in range(n) ) def forward(self, x): ys [self.cv1(x)] ys.extend(m(ys[-1]) for m in self.m) ys.append(self.cv2(x)) return self.cv3(torch.cat(ys, dim1))在YAML层面只需要把原来使用C3k2的位置改成C3k2_RAB并保证模型解析器能识别这个类名。多数YOLO系开源代码的parse_model里都会维护一个“模块类名到构建方式”的分支表比如C2f、C3k2各写一个分支。加入C3k2_RAB以后在分支表里补一个elif m is C3k2_RAB按原来的方式处理参数列表即可。Batch size不敏感输入分辨率512甚至640都能跑不用额外调整。3. 训练配置与实验效果3.1 我自己使用的训练配置我用的是某20类自定义数据集训练集约6万张验证集1500张包含不少低光照和遮挡场景对小目标的考验比较大。训练分辨率统一640x640batch size 32初始学习率0.01使用SGD配合momentum和cosine退火总共训300个epoch。优化器我只试过SGD和AdamW两种最终用的是SGD因为RAB这种轻量模块在SGD下收敛更稳AdamW前期虽然快但到后期mAP的抖动明显更大。这里有个训练策略值得单独说我先冻结backbone权重训练前150个epoch只更新颈部和头部的参数让RAB先在一个相对稳定的特征空间里学习。第150个epoch之后才解开backbone全部参数一起微调。冻结期过早打开backboneRAB产生的新特征分布会干扰预训练权重的统计信息导致Loss长时间卡住。这个体验跟“预训练微调”的通用经验是一致的只是RAB这类注意力模块对特征分布变化更敏感。输入增强上我只用了常用的Mosaic增强第70个epoch以后自动关闭。RAB对几何增强并不敏感所以不需要为了适配注意力模块单独调大rotation或shear参数。3.2 对比基线与其它注意力模块最终结果对比如下所有模型都使用相同的训练配置和验证集单卡训练。模型PRmAP50mAP50-95ParamsGFLOPs基线YOLO260.8120.7430.7690.52311.2M28.4SE0.8050.7510.7720.52811.4M28.7CBAM0.8180.7490.7760.53111.5M29.0RAB本文方案0.8270.7580.7840.53911.3M28.9RAB在mAP50上比CBAM高0.8个百分点mAP50-95高0.8个百分点最明显的变化是召回率R从0.749涨到0.758。这说明注意力机制带来的主要是“把原来漏检的目标找回来”而不是单纯地提升分类置信度。从参数量看RAB只比基线多了0.1M左右远低于CBAM增加的0.3M这是因为CBAM在空间注意力里用了全连接式的压缩而RAB只用了两个小卷积。我还做了消融实验验证RAB内部两个子模块各自的作用只用通道注意力分支mAP50比基线高0.7%只用空间注意力分支mAP50比基线高0.5%两者合并后提升1.5%说明通道注意力与空间注意力有正向叠加效果不是单向压制。3.3 用热力图确认注意力区域是否合理精度指标只是最终结果我习惯在训练到中后期直接看特征热力图确认RAB到底把注意力放在哪里。方法很简单给C3k2_RAB模块注册一个forward hook把输出特征图导出来然后对通道维度做均值归一化之后用colormap渲染到原图上。import cv2 import numpy as np feature_blob {} def hook_fn(name): def hook(module, input, output): feature_blob[name] output.detach() return hook # 注册hook选择某个深层的C3k2_RAB stage_rab model.stage4_c3k2_rab stage_rab.register_forward_hook(hook_fn(rab_feat)) # 推理一张测试图 model.eval() with torch.no_grad(): preds model(test_img) feat feature_blob[rab_feat] # 取第一张图通道维度求均值 heatmap feat.mean(dim1).squeeze(0).cpu().numpy() heatmap (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() 1e-8) heatmap (heatmap * 255).astype(np.uint8) heatmap_color cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) img_vis cv2.addWeighted(test_img_uint8, 0.6, heatmap_color, 0.4, 0)从我的可视化结果看RAB的高响应区域集中在目标主体上背景像素的响应值普遍很低。在遮挡场景下可见的那部分目标区域依然有强响应而普通SE的热力图在半遮挡目标上往往会被切成好几块孤立区域。这一点也解释了为什么RAB对召回率的提升更明显空间注意力分支通过7x7大卷积核建立了区域相关性即使目标的一部分被遮挡模型仍然倾向于把整片目标区域激活出来。4. 训练排障实录4.1 训练初期Loss不降反升第一次复现RAB时我遇到了一个非常明显的现象前5个epoch Loss不降反升从1.2涨到1.8然后才开始下降。排查之后发现两个原因叠加了。第一是RAB初始状态下如果按常规方式初始化为sigmoid(0)0.5恒等路径加注意力路径会把特征幅度放大到原来的约1.5倍深层Stage的BN统计被整体抬高前几个epoch就在修正这个偏差。第二是我当时在外层Bottleneck上也加了RAB浅层特征图分辨率大空间注意力7x7卷积的感受野变小相当于给每个像素做独立判断这会让浅层特征过度锐化。解决办法就是前面提到的负偏置初始化以及只在深层Stage加RAB。改完之后第一个epoch Loss从1.15左右平缓起步后面稳定向下。4.2 注意力权重退化成常数训练到第80个epoch左右我打印了RAB通道注意力的统计值发现均值固定在0.21附近方差几乎为0。也就是说注意力权重退化成常数模块退化成一条固定的线性缩放路径失去了自适应能力。原因有两个。第一是我把reduction设置成4隐藏维度变成128通道压缩太少1x1卷积能直接记住一个常数映射第二是训练时开了超过150个epoch的冻结backbone策略RAB只在这个固定的特征分布里学习提前收敛到局部最优。解决方法是把reduction改回16并且第150个epoch解冻backbone后再把所有RAB模块的学习率调低10倍。这样既能保证注意力在后期继续微调又不会让权重剧烈摆动。4.3 混合精度下出现NaN使用AMP混合精度训练时我在第220个epoch左右遇到一次NaN Loss。排查定位到RAB的空间注意力分支7x7卷积输出经过sigmoid之后值落在非常小的区间比如0.03到0.05之间这些数值在FP16下精度损失很大反向传播时梯度很容易异常。最省事的做法是让RAB的注意力分支单独走FP32不回落到FP16。代码里可以用装饰器或上下文管理器实现torch.cuda.amp.autocast(enabledFalse) def forward_attention(self, x): ch self.channel_att(x.float()) avg_map torch.mean(ch, dim1, keepdimTrue) max_map, _ torch.max(ch, dim1, keepdimTrue) sp self.spatial_att(torch.cat([avg_map, max_map], dim1)) return ch * sp.float() def forward(self, x): identity x out self.forward_attention(x) return identity out注意这里只把注意力分支的输入转成float32原始特征x还是保留原来的精度前向计算量增加很小。如果不想改代码直接关闭AMP里的FP16也可以但训练速度会掉一些。我实测下来还是单独用FP32分支更划算。4.4 替换模块后模型参数解析报错如果你在训练时报“module not found”或“KeyError”大概率是模型构建代码没有注册C3k2_RAB这个类。不同代码库的parse_model结构不太一样但修改思路一致在模块类型判断列表里补一条分支然后按照原来C3k2的参数顺序构建模块。如果你的库版本过旧可能还需要手动把通道数解析逻辑改成ch[f] c2。还有一个容易忽略的点YAML里写[-1, 3, C3k2_RAB, [256, False, 0.5]]时三个参数分别对应c2、c3k、e与构造函数里c1、c2、n、c3k、e的参数顺序并不完全对应。因为第一个参数c1是自动从上一层的输出通道数取出来的所以你只需要保证YAML传参顺序是[c2, c3k, e]即可。5. 沿着RAB还可以继续做的三个方向RAB这种“注意力残差”的结构本身还有很多变体值得尝试。第一个方向是把固定的残差系数变成可学习的参数比如给恒等路径加一个初始值为1的标量权重让网络自己决定“保留多少原始特征、叠加多少注意力特征”。这比现在不加系数的方式灵活一些但需要加一点正则约束防止收敛后系数直接变成0。第二个方向是引入多尺度空间注意力。我目前用的是7x7单核卷积对中大型目标效果不错但对极小的目标7x7可能已经把整张图的上下文都糊在一起了。可以改成并行两个分支一个3x3关注局部一个7x7关注区域上下文最后按通道相加这样对不同尺寸目标的适应能力更强。第三个方向是尝试把RAB放到更深的主干Stage里比如SPPF模块之前的那一层。C3k2本身是“融合多分支然后输出”的结构RAB的残差特性很适合在这种信息汇聚节点上进行注意力调制。我之前只在C3k2内部做过实验没有在SPPF入口做对比最近正打算补上这组实验。我个人现在使用RAB感受最深的一点注意力不是加的越多越好更不是每个Stage都要覆盖。残差注意力虽然能稳定梯度但浅层特征图分辨率高、语义信息弱空间注意力反而容易引入噪声。这个道理其实在所有注意力模块上都通用只是RAB因为自带残差路径容错率更高看起来“怎么加都不出错”真正用的时候还是要克制。把这个模块控制在深层Stage、控制好初始偏置你会在精度和训练稳定性上同时拿到好处。