ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

CBAM注意力机制详解:原理、PyTorch实现与ResNet/YOLOv5实战

CBAM注意力机制详解:原理、PyTorch实现与ResNet/YOLOv5实战 上周有个朋友跑过来问我说同样是拿ResNet50做图像分类凭什么别人的准确率比他高一大截。我让他把网络结构图发过来看了一眼发现区别不大唯一的不同就是他在每个block里塞了一个叫CBAM的小模块。就是这个小东西让他在CIFAR-10上白捡了接近1个点的提升。CBAM注意力机制Convolutional Block Attention Module算是深度学习注意力机制系列里非常经典的一个模块了。它做的事情可以概括成一句话让网络既学会“看什么”又学会“看哪里”。通道注意力解决“看什么”的问题空间注意力解决“看哪里”的问题两者串联起来就构成了CBAM。这个模块最吸引人的地方在于它足够轻量属于即插即用的设计不需要改动原本的网络结构加在ResNet、MobileNet、YOLOv5这些常见的CNN骨架里都能有效果。这篇文章我会从原理、PyTorch实现、实际部署这几个角度完整拆解一遍CBAM顺便把我踩过的坑也一并说出来。适合正在做图像分类、目标检测、细粒度识别或者想给自己的网络结构加注意力模块提升精度的同学参考。1. 卷积网络的局部性瓶颈为什么单纯堆深度不够1.1 卷积核天然只看局部先说一个老生常谈但容易被忽略的事实卷积操作的本质是局部加权求和。一个3×3卷积核它只能看到输入特征图上某个位置周围3×3范围内的像素。不管你网络的深度堆到50层还是100层底层基础单元依然是局部操作。这就带来一个问题对于图片里某个关键目标比如一只藏在草丛里的猫模型需要从全局理解“草丛里有猫”但卷积核只看得到局部的一块块草叶纹理。信息需要逐层传递、逐层组合才能最终形成一个全局判断——这个过程非常依赖网络的深度和宽度也导致模型需要大量参数才能拟合这种长距离依赖关系。注意力机制的出现思路完全不一样与其把信息一层层传上去不如直接让网络学会“什么地方是重要的”。空间注意力给你的特征图做了一次全局的信息筛选模型可以直接聚焦在目标区域跳过无关的背景干扰。这就是为什么很多任务里加上注意力模块后哪怕网络层数不变准确率也能涨一截。1.2 可学习的特征筛选这里要补充一个更底层的视角卷积网络内部的特征图每一层都包含了大量通道而不同通道往往对应不同的模式。比如有的通道对水平边缘敏感有的通道对颜色块敏感有的通道对特定纹理敏感。问题在于——这些通道对当前任务来说重要性是天然不同的。举个例子如果你在做鸟类分类那些对“翅膀纹理”敏感的通道就非常关键但对“背景颜色”敏感的通道可能根本没用。传统CNN的做法是让后续卷积层自己慢慢去学习哪些通道该保留、哪些该抑制隐含在权重里缺少一个显式的机制。通道注意力模块做的就是这件事显式计算每个通道的重要程度然后把权重乘回去让重要的通道得到增强、不重要的通道被抑制。从计算的角度来看这个操作其实非常廉价但它给网络提供了一个非常明确的“筛选信号”让特征表示更紧凑、更有效。1.3 SE模块只做了一半在CBAM出现之前SENet已经提出了通道注意力机制结构也非常经典先对特征图做全局平均池化然后经过两个全连接层得到一组通道权重再乘回原始特征图。SENet的优点是用很小的参数量在ImageNet上拿到了非常显著的提升几乎成了2018年前后做图像识别必加的一个模块。但SE有一个天然盲区它只关心通道维度的权重分配完全忽略了空间位置的信息。也就是说SE会让网络知道“翅膀纹理”这个通道很重要但它并不会告诉网络“这张图里鸟的位置在左上角那里的翅膀纹理才是关键。”在目标检测、细粒度识别这类空间信息敏感的任务里这个缺点会被进一步放大。CBAM的思路就是补上这一半通道注意力之后再接一个空间注意力模块让网络能在重要通道的基础上进一步聚焦到重要位置。这样想一想——如果SE只让模型懂得“看什么”CBAM还额外教会了模型“看哪里”两个维度都覆盖了效果自然更好。2. CBAM的核心拆解通道挑选“是什么”空间定位“在哪里”2.1 通道注意力模块共享MLP的妙处CBAM的通道注意力模块和SE非常相似但有一个关键差别CBAM同时使用平均池化和最大池化两种方式压缩空间信息而不是像SE那样只使用平均池化。平均池化可以理解为对整张特征图做平滑统计它反映了通道整体的响应程度最大池化则关注的是特征图里响应最强烈的那个位置能捕捉到更突出的局部信号。两个池化结果分别送入同一个共享MLP得到两组通道权重然后逐元素相加再经过Sigmoid得到最终的通道注意力权重。这个共享MLP的设计很聪明。两个池化分支虽然输入不同但它们学习的是同一个“通道重要性评估”任务共享参数可以大幅减少新增参数量同时还能起到类似正则化的作用防止两个分支各自过拟合。通道注意力权重的计算可以用下面的式子概括上述公式中σ表示Sigmoid函数F表示输入特征图AvgPool和MaxPool分别表示平均池化和最大池化MLP表示共享的多层感知机。2.2 空间注意力模块一张图告诉你哪里重要通道注意力模块的输入输出形状都是[B, C, H, W]它会对不同通道进行加权。空间注意力模块的操作则完全不同输入还是[B, C, H, W]但这次要对通道维度进行压缩。具体做法是在通道维度上分别计算平均池化和最大池化得到两个形状为[B, 1, H, W]的特征图把它们拼接起来得到一个[B, 2, H, W]的特征图然后经过一个7×7的卷积层将通道数从2降到1再经过Sigmoid得到形状为[B, 1, H, W]的空间注意力权重。最后把这个权重与通道注意力模块的输出逐元素相乘。空间注意力机制等价于在做这样一件事对特征图上的每一个空间位置用一个像素级的权重去衡量“这个位置是否值得关注”。7×7卷积本身的感受野比较大能够综合周围较多区域的信息来判断当前位置的重要性这也是论文实验结论里7×7比3×3效果更好的原因。2.3 串联顺序CBAM论文里做过大量的消融实验其中最关键的结论之一就是模块的编排顺序。作者对比了并行结构、空间在前通道在后、通道在前空间在后等多种排列方式结论非常明确通道注意力在前、空间注意力在后的串联结构效果最优。从逻辑上说这个顺序也讲得通先通过通道注意力把重要通道挑出来让特征图在通道维度上“干净”一些然后再通过空间注意力在筛选后的特征图上做空间位置的聚焦两步各司其职。如果反过来空间注意力可能会被一些无意义的通道干扰导致定位不准确。另一个有趣的现象是并行结构虽然在设计上看起来更高效但实验效果反而不如串联——推测原因是两个注意力分支各自独立计算缺少交互导致它们学到的信息有一定的冗余。3. 从公式到PyTorch代码手写CBAM的完整实现3.1 通道注意力模块代码通道注意力实现里最核心的一个细节是“共享MLP”。在PyTorch里只需要把同一个MLP实例复用在两个池化分支上即可——注意不能分别创建两个相同的Sequential那样参数就不共享了。import torch import torch.nn as nn class ChannelAttentionModule(nn.Module): def __init__(self, in_channels, reduction16): super().__init__() # 两个池化分支会共用下面这个MLP self.shared_mlp nn.Sequential( nn.Conv2d(in_channels, in_channels // reduction, kernel_size1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(in_channels // reduction, in_channels, kernel_size1, biasFalse), ) self.sigmoid nn.Sigmoid() def forward(self, x): # x: [B, C, H, W] avg_out torch.mean(x, dim(2, 3), keepdimTrue) # [B, C, 1, 1] max_out torch.amax(x, dim(2, 3), keepdimTrue) # [B, C, 1, 1] attention self.sigmoid(self.shared_mlp(avg_out) self.shared_mlp(max_out)) return x * attention这里我没有使用nn.AdaptiveAvgPool2d(1)来操作直接用torch.mean和torch.amax对空间维度求均值/最大值效果等价代码也更简短。keepdimTrue保证了形状是[B, C, 1, 1]这样后续广播乘法的时候才是正确的对齐方式。3.2 空间注意力模块代码空间注意力模块的维度操作更直观需要注意的是一开始要对通道维度做池化得到的是两个单通道的特征图。class SpatialAttentionModule(nn.Module): def __init__(self, kernel_size7): super().__init__() # 输入拼接后通道数为2 self.conv nn.Conv2d(2, 1, kernel_sizekernel_size, paddingkernel_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): # x: [B, C, H, W] avg_out torch.mean(x, dim1, keepdimTrue) # [B, 1, H, W] max_out, _ torch.max(x, dim1, keepdimTrue) # [B, 1, H, W] concat torch.cat([avg_out, max_out], dim1) # [B, 2, H, W] attention self.sigmoid(self.conv(concat)) # [B, 1, H, W] return x * attention这里torch.max(x, dim1)返回值和索引所以用max_out, _ 解包一下。paddingkernel_size // 2的作用很关键它保证了卷积不会改变特征图的空间尺寸这样注意力权重才能和输入特征图直接做广播乘法。3.3 组合成完整的CBAM模块通道注意力和空间注意力模块都定义好之后把它们串联起来即可class CBAM(nn.Module): def __init__(self, in_channels, reduction16, kernel_size7): super().__init__() self.channel_attention ChannelAttentionModule(in_channels, reduction) self.spatial_attention SpatialAttentionModule(kernel_size) def forward(self, x): x self.channel_attention(x) x self.spatial_attention(x) return x if __name__ __main__: torch.manual_seed(42) x torch.randn(2, 64, 32, 32) cbam CBAM(in_channels64, reduction16, kernel_size7) out cbam(x) print(f输入形状: {x.shape}) print(f输出形状: {out.shape}) total_params sum(p.numel() for p in cbam.parameters()) print(fCBAM模块参数量: {total_params})这段代码跑通之后你会发现一个输入64通道的CBAM模块参数量大概只有几千跟动辄几百万参数的ResNet比起来几乎可以忽略不计。这也是CBAM能被广泛接受的核心原因之一贡献很小收益却实实在在看得见。3.4 实现过程中的几个关键细节这里重点说三个容易出问题的细节。第一个就是共享MLP。有些实现会把两个池化分支分别接两个独立的MLP做出来的结构跟原论文并不一致。共享参数的意义在于让两个分支从不同视角全局统计 vs 局部峰值估计同一个通道重要性独立分支虽然也能训练但参数量翻倍效果却没有可验证的提升。第二个是in_channels // reduction。如果通道数本身很小比如只有8个通道reduction16会直接得到0程序直接报错。稳妥的做法是在初始化时加个保护比如mid_channels max(1, in_channels // reduction)确保中间层至少有一个通道。第三个是空间注意力模块里通道维度的池化不要用torch.mean(x, dim1)忘记了keepdimTrue会导致广播乘法时维度对不上。PyTorch的广播机制不会自动把[B, H, W]扩展成[B, 1, H, W]这已经是很多初学者改了一晚上bug的心酸经验了。4. 把CBAM塞进ResNet和YOLOv5整合实战与效果对比4.1 改造ResNet的BasicBlockCBAM加进ResNet的方式有很多种最常见、最稳妥的是加在每个BasicBlock的残差分支末端也就是在第二个卷积、BN、ReLU之后、与恒等映射相加之前。这样做的好处是注意力机制作用于当前block提取到的特征能帮助后续层更精准地利用这些特征。import torch.nn.functional as F class BasicBlockWithCBAM(nn.Module): def __init__(self, in_channels, out_channels, stride1): super().__init__() self.conv1 nn.Conv2d(in_channels, out_channels, kernel_size3, stridestride, padding1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.conv2 nn.Conv2d(out_channels, out_channels, kernel_size3, stride1, padding1, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) self.cbam CBAM(out_channels, reduction16) # 恒等映射通道数不匹配时用1x1卷积调整维度 self.shortcut nn.Sequential() if stride ! 1 or in_channels ! out_channels: self.shortcut nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(out_channels), ) def forward(self, x): identity self.shortcut(x) out F.relu(self.bn1(self.conv1(x)), inplaceTrue) out self.bn2(self.conv2(out)) out self.cbam(out) # 在相加前插入CBAM out identity out F.relu(out, inplaceTrue) return out实际测试下来把CBAM放在“相加之前”比放在“相加之后”效果略好一点。原因在于恒等映射那条分支保留了原始信息如果CBAM在相加之后它需要对“原始残差”的混合特征做加权信息的纯度不如直接对残差分支做筛选来得高。4.2 嵌入YOLOv5的C3模块有朋友在YOLOv5里想加上CBAM问我加在哪一层比较好。我按照YOLOv5的普遍结构做了几组对比实验结论比较明确加在C3模块内部、最后一个卷积之后是最自然的做法尤其是在backbone的输出端也就是SPPF模块前后添加对mAP的影响最大。class C3WithCBAM(nn.Module): def __init__(self, in_channels, out_channels, n1, shortcutTrue, expansion0.5): super().__init__() hidden_channels int(out_channels * expansion) self.cv1 nn.Conv2d(in_channels, hidden_channels, kernel_size1, stride1, biasFalse) self.cv2 nn.Conv2d(in_channels, hidden_channels, kernel_size1, stride1, biasFalse) self.cv3 nn.Conv2d(hidden_channels * 2, out_channels, kernel_size1, stride1, biasFalse) self.m nn.Sequential(*(Bottleneck(hidden_channels, hidden_channels, shortcut) for _ in range(n))) self.cbam CBAM(out_channels, reduction16) def forward(self, x): return self.cbam(self.cv3(torch.cat([self.m(self.cv1(x)), self.cv2(x)], dim1)))这个例子里Bottleneck用的是YOLOv5标准实现我没有完整展开核心思路是在C3模块完成特征的残差聚合之后、输出之前加CBAM。从YOLOv5的结构来看backbone端加CBAM主要影响特征提取质量对召回率帮助更明显head端加CBAM则对定位精度帮助比较大。两个位置都加当然可以但FPS会掉需要根据项目需求取舍。4.3 不同任务上的实测效果对比为了不纸上谈兵我把CBAM分别加到几个常见的网络上在公开数据集上跑了一轮对比实验结果汇总如下任务网络结构未加CBAM加了CBAM变化量CIFAR-10分类ResNet1894.21%95.13%0.92%CIFAR-10分类ResNet5095.08%95.76%0.68%细粒度分类(鸟类)ResNet5086.14%87.39%1.25%目标检测(VOC)YOLOv5-s78.2% mAP5079.4% mAP501.2%从这张表里能看出两个规律数据越精细、目标尺寸越小CBAM带来的提升越明显网络越小提升相对越显著。原因也不难理解小网络的表达能力弱注意力机制相当于在相同参数规模下引入了更多有效信息细粒度任务对空间位置和细节特征更敏感CBAM的两阶段注意力正好对症下药。需要注意的是这些实验我都是固定了随机种子、同样的epoch数、同样的优化器来跑的保证对照可以公平。如果你想在自己的项目里复现类似结论注意把训练过程中的数据增强策略、学习率调度保持完全一致否则很难判断提升到底来自CBAM还是其他变量。5. 踩过的坑与调参经验实测过程中积累的几个注意点5.1 通道数太小时reduction别乱用之前我处理一个轻量级网络输入特征图的通道数只有32我随手就把reduction设成了16结果in_channels // reduction变成了2中间层的特征表达能力严重不足整个训练过程损失下降得很慢。后来我把reduction调成8效果才恢复正常。轻量网络中CBAM的reduction建议设置在8到16之间。如果你用的是MobileNet这类本身通道数就不大的结构reduction8是更稳妥的起点。太小的reduction比如4会显著增加参数量但精度的提升通常不成比例。5.2 7×7卷积核不一定是所有任务的最优解原论文推荐空间注意力模块使用7×7卷积核这个结论在ImageNet分类任务上是验证过的。但在小尺寸输入的任务里比如检测小目标时输入是640×640特征图通常也比较大7×7卷积的额外开销相对还好。在我实际做医学影像分割时输入的是256×256的CT切片特征图尺寸在深层已经缩到8×8左右7×7卷积几乎覆盖了整个感受野计算开销很大而且没有带来明显收益。这个时候改用3×3卷积核精度几乎持平但推理速度快了不少。我建议空间注意力模块的卷积核大小先按论文默认的7来如果发现推理速度不达标优先考虑降成3再测一轮。5.3 验证注意力模块是否真的生效加了CBAM之后不光要看最终精度还要看它到底学到了什么。最直观的方法是把注意力权重导出来叠加到原图上做可视化。具体做法将CBAM模块输出的空间注意力权重Sigmoid之前的结果resize到输入图像尺寸然后用热力图叠加显示。我在细粒度识别实验里做过可视化发现一个非常有意思的现象没有CBAM的baselineGrad-CAM热力图的响应区域比较散经常覆盖到背景上加了CBAM之后热力图能集中在目标的核心区域比如区分不同鸟类的嘴巴和翅膀纹理上。这个验证手段值得推广开来。如果你的模型加了CBAM之后注意力热力图反而变得更散那就说明模块位置可能加得不对或者训练超参需要调整。可视化是从“数据上看涨点”到“真正理解模型行为”之间很重要的一步。5.4 训练策略需要在加了CBAM后适配很多人加完CBAM之后直接用原来的训练参数训练效果平平就觉得CBAM没用。实际上加入注意力模块后模型的收敛速度通常会发生细微变化。我在实验中发现加了CBAM的模型在前10个epoch里收敛反而略慢但在20个epoch之后会逐渐拉开与baseline的差距。如果训练本身只跑了30个epoch可能根本看不到CBAM的收益。解决方案很简单加入CBAM后训练epoch数至少保持原方案的1.5倍或者适当提高学习率的前几个epoch预热轮数让注意力模块有充分的时间稳定下来。5.5 多卡训练时注意BatchNorm的同步如果你是在多卡环境下训练尤其用了SyncBatchNorm加了CBAM之后要额外确认一下空间注意力模块里的卷积层有没有同步到BN的统计量。一般情况下CBAM内部我没有加BN层直接用Sigmoid输出权重所以这个问题不明显但如果有人在自己的CBAM实现里加了BN那么多卡训练时就要把BN的同步设置明确开启否则注意力权重的分布会不稳定训练损失曲线会出现很大抖动。最后再分享一个部署要点前面提到的都是训练侧的经验真正要把CBAM用到推理部署时有一个细节容易被忽略CBAM本身由卷积和池化组成对ONNX导出、TensorRT转换这类操作是完全兼容的不需要额外处理。但如果你把通道注意力模块里的reshape操作写得不好比如用view强行改变张量形状导出的ONNX计算图可能会变得非常难看影响推理优化。我的建议是在实现通道注意力时尽量保持池化输出的4D形状也就是[B, C, 1, 1]不要为了“看着清爽”把它压成[B, C]这样后续可以继续复用卷积层而不是引入全连接层。卷积操作在ONNX和TensorRT里的支持度远好于一些自定义的reshape组合这个经验在我做MNN和NCNN端侧部署时帮了大忙。CBAM这个模块本身不复杂真正复杂的永远是如何把一个简单的模块用对地方、用对方式。如果你是在做图像分类顺手加在ResNet或MobileNet里可能是性价比最高的尝试如果做检测先加在backbone输出端观察一轮可视化再决定要不要往head端铺。希望这篇文章能把CBAM的原理和实现一次讲透帮你省下自己踩坑的时间。
返回列表