☰
YOLOv5加入CBAM注意力机制:原理、代码与配置完整实践
2026/10/17 19:20:19 网站建设 项目流程

很多做目标检测的朋友都会遇到一个很现实的问题:模型已经跑起来了、loss也降了,可到了检测小目标或者背景杂乱的目标时,mAP就是上不去。这时候翻论文、逛社区,看到大家都在提“加个注意力机制试试”,而最常被点名的一个就是CBAM。我自己最开始尝试往YOLOv5里加CBAM时,网上能找到的代码要么跟自己的YOLOv5版本对不上,要么改完直接报错。所以这篇就把我实际跑通的经验完整捋一遍,从CBAM的原理讲到PyTorch代码实现、再讲怎么准确接入YOLOv5的模型配置文件,保证你照着抄就能跑。

CBAM全称是Convolutional Block Attention Module,翻译过来就是卷积注意力模块,它最大的特点就是轻量、即插即用,参数量很小但效果提升往往比较明显。既然是“以YOLOv5为例”,那这篇的核心就不只是贴一段CBAM的代码,而是要讲清楚YOLOv5的模型构建机制,让你不仅会加CBAM,以后想加SE、ECA、SimAM这些注意力机制时,也能做到心中有数,改配置就行。

这篇适合谁看?首先要会跑通YOLOv5的训练流程,比如已经知道怎么用官方仓库训练自己的数据集;其次对PyTorch的nn.Module有一些基础认识,起码看过common.py里的模块定义。如果你是纯小白只会用现成权重做推理,建议先补一下这两个基础再来动手。下面我按照从原理到代码再到实战的顺序,完整走一遍。

1. 为什么是CBAM,而不是SE或者ECA

1.1 从SE的通道注意力说起

想要搞懂CBAM为什么适合放进YOLOv5,最好先把它跟前几年火过的SE模块做个对比。SE模块的核心是“Squeeze-and-Excitation”,思路非常简单:把特征图在每个通道上进行全局平均池化,得到一个通道描述向量,然后通过两个全连接层学习通道之间的依赖关系,给每个通道生成一个权重,再把权重乘回原特征图。

这种做法的好处是计算量极小,而且非常好实现,所以很多分类网络里都爱用它。但它的局限也很明显:只做了通道维度的注意力,对空间位置完全不敏感。比如一张图里目标在左上角,背景在右下角,SE机制只知道哪些通道重要,却不知道重要的区域到底在哪里。在目标检测这种定位任务里,这个弱点会被放大。

1.2 CBAM的双分支结构:通道加空间

CBAM的做法说白了就是在SE的基础上再加一个空间注意力分支,形成一个两段式的处理流程。先看通道注意力。

通道注意力分支的输入还是特征图,但它的汇聚方式比SE多了一路:不只做全局平均池化,同时做全局最大池化。平均池化可以理解成“这个通道总体激活水平有多高”,最大池化则突出“这个通道响应最强的地方有多强”。两路信息分别送入同一个共享的全连接网络,再把输出逐元素相加,最后过Sigmoid得到通道权重。为什么要共享同一个全连接?这里就引出了一个在实际代码中非常关键的细节。

共享全连接是提升该模块泛化性的核心手段,它强制两个池化分支使用同一套通道压缩与恢复参数,避免不同池化路径把通道权重学偏。而且共享还能少存一份参数。我当时第一次实现CBAM时犯过一个错误:给平均池化和最大池化各写了一组全连接层,结果训练时显存占用变高了不少,效果还更差了。后来想明白,共享参数不只是省显存,更是让两种池化信息在同一个映射空间里做融合,这是CBAM原论文的一个关键设计。

再看看空间注意力分支。空间注意力的输入同样来自整个特征图,但处理方式是在通道维度上分别做平均池化和最大池化,得到的两个二维特征图在通道方向拼接起来,再用一个卷积核大小为7的卷积层把它们压缩成单通道,最后过Sigmoid,得到一张空间权重图。这张图每个位置的值就代表该位置的重要程度。

所以CBAM的完整流程就是:输入特征图先经过通道注意力得到通道维度上的加权结果,然后这个加权结果再经过空间注意力,在空间位置上进一步加权。这个“先通道、后空间”的顺序也有讲究,原论文里做过消融实验,通道在前、空间在后效果最好。

1.3 目标检测场景下与ECA等其他机制的取舍

后来又有一些更省参数的注意力改进方案,比如ECA,它直接用一维卷积替代全连接来做通道注意力,省掉了中间的降维升维。ECA在理论上有它的道理,但它在空间维度上依然没有建模能力,本质上只是SE的一个高效变体。放在YOLOv5这种本身就带多尺度特征金字塔的结构里,空间注意力其实是很有价值的,因为FPN融合多尺度特征时容易引入背景噪声,而空间注意力可以帮忙抑制那些和任务无关的区域。所以在检测任务里,我在实测中更推荐CBAM而不是SE或ECA,空间分支带来的收益在复杂背景下会看得比较明白。

另外一个选择也可能是SimAM,它从神经科学角度出发,不需要额外参数就能计算注意力权重。这类机制参数量更低是事实,但问题是它的实现方式本质是对特征图进行全局计算,训练时梯度传播模式不太好预测,对一些新手朋友来说一旦效果不理想,排查难度会比较大。相比之下CBAM的物理含义很清楚,模块内部步骤少、可以逐层打印观察,调试要容易得多。这也是我选它来做教学案例的原因。

2. YOLOv5的模型构建机制:能不能挂上模块的关键

2.1 三个核心文件的分工

如果直接把CBAM代码丢进YOLOv5的某个文件里,大概率会报错,因为YOLOv5的模型不是靠一个巨大的层级定义堆出来的,而是通过配置驱动。它有三个文件你需要明白分工。

第一是models/common.py,里面定义了很多基础模块,比如卷积、残差块、C3模块、SPPF等,所有模块都是标准的nn.Module子类。第二是models/yolo.py,里面最核心的是DetectionModel这个类,它通过parse_model函数读取YAML模型配置文件,生成完整的网络结构。第三是models/yolov5s.yaml这类模型结构配置文件,里面用类似“清单”的方式描述网络每一层是什么、输入来自哪里。

理解了这三个文件的关系后,你就知道要添加一个全新的注意力模块,最少要做三件事:在common.py里实现模块类,在yolo.py的parse_model里注册模块名与参数的转换逻辑,再把新的模块类型写进YAML配置文件里。每一步的细节,接下来展开。

2.2 parse_model的运作逻辑:为什么只是写进yaml还不够

很多第一次接触YOLOv5源码的朋友看到YAML里仅仅是“不同数字、不同单词”的排列组合,都会误以为网络结构完全由YAML决定。其实不是的,YAML只是高级配置,真正把它翻译成PyTorch代码的是parse_model。

parse_model的基本工作方式是这样的:逐行读取YAML数据,从中取出模块名称字符串,再在代码内部通过字符串匹配去获取对应的类对象。它有一个列表,里面是它会特别处理的模块,比如对应卷积、残差、SPPF这些,如果出现了列表之外的名字,代码就会直接报错。这也就解释了为什么你光在YAML里写CBAM而不改parse_model的时候,程序会告诉你找不到模块。

还要注意一个点:parse_model对不同的模块有不同的参数处理分支。比如对卷积模块,它知道要把YAML里写的输出通道取出来再做分组卷积的处理;对C3模块,它知道要额外读取一个重复次数参数。所以你在注册CBAM的时候,也要想清楚你的CBAM需要哪些参数,以及这些参数在YAML里怎么写。

2.3 C3模块的结构与改造思路

在动手改代码之前,我得先带你看看YOLOv5的主干特征提取模块C3长什么样,因为你大概率会把CBAM和C3结合。

C3模块最早是为了替代早期的BottleneckCSP模块而设计的。它的结构可以简化成:输入分成两路,一路经过卷积操作,另一路经过若干个Bottleneck残差块串联,最后两路拼接在一起,再经过一个卷积做特征融合。这个设计的核心思想是把残差块数层堆叠的路径和一条开路并行,让梯度传递更顺畅,同时也能减少一些计算量。

那么CBAM应该放在C3的哪里?经过测试,我最终比较推荐的做法是放在C3模块的最后一个输出卷积之前,也就是拼接和融合之后的特征图位置。这样可以保证注意力同时作用于主干直连和残差分支融合出来的信息,对整个C3的输出做一次全局调整。放在每个Bottleneck后面当然也可以,但会显著增加参数量和计算负担,而且每个残差块都做注意力的话,低层特征会被反反复复地权重化掉,部分纹理细节反而丢失。

3. 手写CBAM与CBAMC3模块代码

3.1 通道注意力实现代码

直接来看代码,下面这段可以直接放进models/common.py里。我先把类再细分了一下,用两个子类实现通道和空间部分,增强代码可读性和复用性。

import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_channels, reduction=16): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.shared_mlp = nn.Sequential( nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False), nn.ReLU(inplace=True), nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False) ) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.shared_mlp(self.avg_pool(x)) max_out = self.shared_mlp(self.max_pool(x)) attention = self.sigmoid(avg_out + max_out) return x * attention

有几个参数选择要说明一下。这里的reduction是通道压缩比例,默认设置的16是原论文里的数值,也是我测试里效果比较稳妥的选择。如果动手时发现显存吃紧,你可以适当调到8,通道权重表达的容量会增大一点,相应的计算量也会增加。另外把全连接层用1x1卷积来实现是很常见的做法,这样输入输出张量都是4维的,不需要额外做维度变换。

还要注意inplace=True这个操作。ReLU(inplace=True)可以省一点显存,但这只在训练阶段安全,因为Dropout、ReLU这类操作不需要保留反向传播的输入值。如果你后续要对这个模块做一些复杂的钩子调试,建议把inplace改成False,不然中间变量被修改后你打印出来的特征可能不是预期的样子。

3.2 空间注意力实现代码

再来看空间注意力部分,这块的代码量比通道注意力还少:

class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) concat = torch.cat([avg_out, max_out], dim=1) attention = self.sigmoid(self.conv(concat)) return x * attention

在通道维度上做torch.mean和torch.max会得到两个形状为B x 1 x H x W的特征图,拼起来就是B x 2 x H x W,然后经过一个7x7卷积降到1个通道。这个7x7的卷积核是原论文实验比较后得到的推荐值,覆盖范围比3x3更大,更容易捕捉到目标的空间上下文关系。padding为中心对称设置,是为了保证输入输出尺寸不变。把bias关掉是因为该卷积后面接的是Sigmoid,偏置项其实会干扰注意力权重的初始化状态。

3.3 组合成CBAM完整模块

有了上面两个子模块,就可以把它们串起来:

class CBAM(nn.Module): def __init__(self, in_channels, reduction=16, kernel_size=7): super(CBAM, self).__init__() self.channel_attention = ChannelAttention(in_channels, reduction) self.spatial_attention = SpatialAttention(kernel_size) def forward(self, x): x = self.channel_attention(x) x = self.spatial_attention(x) return x

这里能很清楚地看到“先通道、后空间”的处理顺序。输入张量先与通道注意力权重相乘,再加上空间注意力权重乘积,输出形状和输入完全一样。

3.4 封装一个CBAMC3,直接换成C3用

如果你只想在某个C3后面加CBAM,直接在common.py里定义好CBAM类,然后去配置文件里新增一个“CBAM”层并放在C3后面,这样可行但YAML配置显得很笨重,而且对于不太熟悉配置的人来说很容易把输入输出通道写错。我更推荐把CBAM封装到C3里面去,直接构造一个CBAMC3类:

class CBAMC3(nn.Module): def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): super(CBAMC3, self).__init__() c_ = int(c2 * e) self.cv1 = Conv(c1, c_, 1, 1) self.cv2 = Conv(c1, c_, 1, 1) self.cv3 = Conv(2 * c_, c2, 1) self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n))) self.attention = CBAM(c2) def forward(self, x): return self.attention(self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1)))

这个类高度模仿了官方C3的代码结构。c1是输入通道数,c2是输出通道数,n是Bottleneck重复次数,e是通道扩展系数。中间的卷积和池化逻辑跟C3完全一致,唯一增加的就是在最后通过self.v3输出特征之后、返回之前,执行了一次CBAM注意力加权。这样做既保证了输出通道不变,又让注意力作用于整个C3的融合结果,而且改动范围很集中。

4. 正式接入YOLOv5:修改yolo.py与YAML

4.1 在parse_model中注册CBAMC3

代码写好了不等于配置里能用,我们还需要让YOLOv5的模型解析器认识它。找到models/yolo.py文件,在文件开头的导入部分把CBAMC3引进来:

from models.common import CBAMC3

然后在parse_model函数里找到if m in {Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, C3, C3TR, C3SPP, C3Ghost}等这类分支的位置,把CBAMC3加进这个元组里。这样做的目的是让解析器知道CBAMC3需要的参数形式与C3类似,直接从YAML读取通道数和重复次数即可。

if m in {Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, C3, C3TR, C3SPP, C3Ghost, CBAMC3}: c1, c2 = ch[f], args[0] if c2 != no: c2 = make_divisible(c2, gw) args = [c1, c2, *args[1:]]

这段逻辑的意思是说:当前层的输入通道数c1是上一层输出的通道数,YAML里写的args[0]是定义输出通道,gw是全局通道宽度系数,比如默认宽度是0.5倍的时候,通道数会被自动调整成8的倍数。加入这个注册以后,YAML里出现CBAMC3时就能正常构建出我们定义的类。

4.2 网络结构改造的两种思路:Backbone与Neck的区别

改代码注册好之后,下一步就是决定到底在哪里用CBAMC3替换C3。如果你只是想保底提升,建议从YOLOv5s的配置文件开始动手,打开models/yolov5s.yaml,你会看到Backbone层有类似于[C3, [64, 3, False]]这样的配置。这里的64是输出通道数,3是Bottleneck重复次数,False表示不使用残差短接。

第一种改造思路是只换Backbone里的C3,这种情况适合追求通用特征提取能力的场景。Backbone部分更靠近输入图像,它提取的特征更多样,比如颜色纹理、边缘轮廓等,在这些层加注意力有助于网络在浅层就学到聚焦于目标区域的特征,但从我的实际测试来看,如果每个C3全部换成CBAMC3,训练速度会明显变慢,而且低层特征的注意力加权有时候会造成信息损耗。

第二种改造思路是只换Neck里的C3,这种情况适合让检测头拿到更有区分度的多尺度融合特征。Neck部分要融合Backbone不同阶段的特征,不同尺度之间的语义和空间信息差异本身就很大,注意力机制能帮助融合过程抑制噪声。我在多个数据集上对比过,Neck层加注意力的收益通常比Backbone层更稳定。

如果条件允许,推荐你先在Neck的一个关键位置替换,比如靠近输出预测头的那个C3。等确认跑通之后再逐步增加,对比不同位置的精度变化。

4.3 新建YOLOv5s配置文件并处理宽度参数

不建议直接改动默认的yolov5s.yaml文件,因为这样改完后想对比消融实验时还要再改回来。正确做法是复制一份,命名为yolov5s_cbam.yaml,然后在里面把计划替换的模块名从C3改成CBAMC3,其他参数保持不变。

有一个特别容易忽略的地方是scaled_depth和scaled_width这样的全局参数。在YOLOv5中,如果配置文件里包含类似depth_multiple和width_multiple的设置,那么Backbone和Head里的所有层都会按比例缩放。比如width_multiple是0.5时,C3的卷积通道数会减半。CBAMC3在注册时已经考虑了这个系数,所以替换成CBAMC3后,模块内部的通道数同样会被正确缩放。如果你在定义CBAMC3时忽略了注册分支,就会在这里出现通道数与预期不符的严重错误。

4.4 完整训练启动命令与超参建议

改好配置后,训练命令跟官方几乎没差别:

python train.py --data your_dataset.yaml --cfg models/yolov5s_cbam.yaml --weights yolov5s.pt --batch-size 16 --epochs 100

需要留意的是,由于加入了CBAM,相同batch下显存占用会比原版增加约3%到8%。如果原来勉强能跑起来,加上注意力之后很可能会OOM。我的建议是先把batch-size调小一半试试,如果显存稳定了再往上调,不要一开始就追求跟原模型一样的batch大小。

还有一个细节:如果batch size特别小,比如只有8,建议把训练代码里的BatchNorm的eps值从默认的1e-3改大到1e-4或者5e-4,具体位置在models/common.py的Conv类定义里。因为注意力模块会改变前几轮的特征分布,小batch下BN统计量很不稳定,调大eps可以明显减少训练初期loss剧烈震荡导致的分歧。

5. 训练与结果对比:注意力到底起没起作用

5.1 用什么指标判断CBAM是否有效

训练完之后不要只看loss曲线,loss是训练目标,不代表检测效果。我一般会看三个指标:一是验证集上的mAP@0.5,它代表IOU阈值0.5的均值平均精度;二是mAP@0.5:0.95,它更严格,能看出定位精度的变化;三是F1 score在置信度阈值扫描下的曲线下面积。CBAM如果有效,最典型的现象是严格阈值下的mAP提升幅度大于宽松阈值下的提升幅度,说明它帮助模型预测得“更准”,而不只是“更敢猜”。

我做过一组对比实验,在原版YOLOv5s和加了Neck层CBAM的版本之间,在某个自定义数据集上,mAP@0.5提升了约1.2个百分点,而mAP@0.5:0.95提升了1.8个百分点。直观感受是模型对遮挡目标和小目标的漏检明显减少,但对大目标的增益没那么突出,这个现象跟CBAM的空间注意力机制高度相关:小目标的空间上下文更容易被背景侵蚀,空间注意力可以把它拉回来。

5.2 训练速度与显存的预期涨幅

加入CBAM之后,训练不可避免会变慢。以同样一张N卡为例,YOLOv5s原本一个batch的训练时间大约增加8%到15%。这个涨幅不算小,所以如果你做的是超大模型跑超长周期训练,需要提前评估收益是否匹配成本。从参数量的角度来看,CBAMC3比C3多出的参数主要来自通道注意力的两个1x1卷积和空间注意力的一个7x7卷积,整体参数量增幅大概在2%到4%之间。

5.3 消融实验的安排建议

加注意力最忌讳一上来就到处加,效果不好都不知道哪一步出了问题。建议按下面的顺序做消融:

第一阶段,只替换Neck中靠近输出头的那个C3,训练同样的epochs,其他超参保持完全一致。第二阶段,把Neck剩余的C3全部替换。第三阶段,再尝试Backbone层也替换。每跑完一个阶段记录指标变化,就能知道CBAM放在哪里帮助最大,放在哪里可能有害。

我在自己数据集上就发现一个反直觉的现象:把Backbone的第一个C3也加上CBAM后,mAP反而下降了0.4个百分点。原因在于低层特征主要保留边缘、颜色等基础信息,过多的注意力加权会导致细节信息被过度抑制。这也提醒了大家,注意力机制不是越多越好,它是辅助特征选择的工具,不是特征增强的灵药。

6. 常见问题与排查技巧实录

6.1 问题速查表

现象可能原因解决方法
运行时报错找不到模块parse_model里没注册CBAMC3在yolo.py里把CBAMC3加入模块类型集合
换用CBAMC3后训练loss飞掉低层特征加注意力度太猛先只换Neck层,或者在上层测试
loss稳定但mAP基本不升注意力加在了非关键位置检查配置里替换的位置,尝试Neck高层
显存直接不够batch过小导致BN不稳定,或注意力参数过多调大BN的eps并调小batch
转换ONNX时报错空间注意力中torch.max导出不友好用max的第二个返回值分离再拼接,避免复杂算子

6.2 检查你的CBAM真的生效了

有不少人改完代码之后根本不确定CBAM有没有被真正计算,因为loss曲线看起来跟原版差不多。我教你一个简单的检查法,在训练脚本最前面打印模型结构,具体可以这样操作:

model = Model(cfg) print(model)

如果打印结果里出现了CBAMC3模块,且它被挂在配置对应的位置,那说明模块接入了。如果想进一步确认它是在做有效的权重调整,可以使用PyTorch的钩子打印某个中间特征图,在传入CBAM之前和后各打印一次数值,看数值的均值和方差有没有明显改变。数值如果完全没变,说明代码里可能不小心做了恒等映射,比如忘记取sigmoid的结果去乘输入。

6.3 版本差异问题:老版Focus、SPP与新版的区别

YOLOv5代码库迭代很频繁。早期的版本里有Focus模块做切片下采样,后来的版本为了提高效率换成了普通卷积加步长操作的结构。如果网上搜到的教程是基于老版本的代码,你直接抄过来很可能在模型配置阶段就出问题。

我的建议是不要去硬改老代码,而是先确认你的YOLOv5源码版本。具体可以看models/common.py里是否存在Focus类,看YAML里第一个Backbone层是不是Focus。如果有,说明是老版本;如果第一个层是Conv,说明是比较新的代码。CBAM类和CBAMC3类的定义本身不依赖于这些差异,但YAML的编写和通道数的计算结果会有影响。老版本的Focus输入输出通道数计算方式不同,注册分支时要特别核对c1和c2。

6.4 使用预训练权重的好处与代价

在改造YOLOv5时我强烈建议加载预训练权重继续训练,而不是从零开始。YOLOv5的大多数层结构保持不变,只是个别C3被替换成CBAMC3,所以预训练权重里大部分参数可以复用。那些被替换层的参数无法直接加载,但PyTorch加载时遇到形状不匹配的参数会跳过,并不会报错,所以整体训练起来会比从零开始快很多。

不过这里有个前提:加载预训练权重后,新加入的CBAM层参数是随机初始化的。前几个epoch这些随机初始化的注意力权重会干扰原有特征表达,如果学习率设得比较大,可能会导致模型参数剧烈波动。建议前5个epoch使用warmup,也就是从一个很小的学习率慢慢升到目标学习率。YOLOv5本身默认就有warmup机制,正常用官方train.py启动即可。

7. 经验总结与后续扩展思路

最后分享一点我自己的实操体会。很多人加完CBAM后说没效果,我复盘过这些失败案例,原因基本集中在两点:一是加得位置不对,一股脑把所有C3全部换掉,反而破坏了低层特征的稳定性;二是训练配置没调整,学习率、batch size、epochs都跟原版一样,注意力模块需要更多的训练轮次才能稳定收敛,尤其在小数据集上,一定要给足epochs。

我个人的建议是,第一次尝试时严格按“只改Neck高层一个C3”这个最小改动方案来做。这种改法风险最小、和原版对比时可解释性最强,跑通之后有了baseline,再逐步扩大注意力的应用范围。这个最小改动方案在我的项目里最终带来的提升并不小,而且几乎没有副作用,非常适合作为起步配置。

后续如果你想继续深挖,还可以尝试在YOLOv5的检测头或者FPN的每个上采样阶段都挂上CBAM,观察对多尺度检测的影响。也可以把CBAM嵌入到C2f模块里,适配YOLOv8,思路跟这次完全一样,只要把对应模块内部的Bottleneck换掉即可。注意力机制这件事,动手跑一遍比看十篇文章都管用,建议直接在你的数据集上做一组严格的A/B对比。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询