YOLOv5改进系列(7)——添加SimAM注意力机制
🚀一、SimAM介绍
- 论文题目:《SimAM: A Simple, Parameter-Free Attention Module for Convolutional Neural Networks》
- 论文地址:http://proceedings.mlr.press/v139/yang21o/yang21o.pdf
- 代码地址:GitHub - ZjjConan/SimAM: The official pytorch implemention of our ICML paper "SimAM: A Simple, Parameter-Free Attention Module for Convolutional Neural Networks".
1.1 简介
SimAM (Simple Attention Mechanism) 是中山大学在注意力机制方面的尝试,从神经科学理论出发,构建了一种能量函数挖掘神经元重要性,并对此推导出了解析解以加速计算。通过ImageNet分类、COCO检测与分割等任务验证了所提SimAM的灵活性与有效性。值得一提的是,所提SimAM是一种无参数注意力模块。
SimAM的设计思路源于SENet,但不同于SENet的复杂结构,SimAM只使用了一个全局池化层和几个全连接层。具体来说,我们在YOLOV5的倒数第二个卷积层后面加上一个全局池化层,将卷积层的输出向量变为标量。然后,我们将这个标量输入到两个全连接层中,以得到两个权重向量。这两个权重向量分别用于加权卷积层的特征图和辅助分类器的特征图。
1.2 方法
之前的方法
- (a)通道注意力:1D注意力,它对不同通道区别对待,对所有位置同等对待;
- (b)空间注意力:2D注意力,它对不同位置区别对待,对所有通道同等对待。
本文的方法
权值生成方法。现有注意力往往采用额外的子网络生成注意力权值,比如SE的GAP+FC+ReLU+FC+Sigmoid。更多注意力模块的操作、参数量可参考下表。总而言之,现有注意力的结构设计需要大量的工程性实验。我们认为:注意力机制的实现应当通过神经科学中的某些统一原则引导设计。
🚀二、在backbone末端添加SimAM注意力机制方法
2.1 添加顺序
(1)models/common.py --> 加入新增的网络结构
(2) models/yolo.py --> 设定网络结构的传参细节,将SimAM类名加入其中。(当新的自定义模块中存在输入输出维度时,要使用qw调整输出维度)
(3) models/yolov5*.yaml --> 新建一个文件夹,如yolov5s_SimAM.yaml,修改现有模型结构配置文件。(当引入新的层时,要修改后续的结构中的from参数)
(4) train.py --> 修改‘--cfg’默认参数,训练时指定模型结构配置文件
2.2 具体添加步骤
第①步:在common.py中添加SimAM模块
将下面的SimAM代码复制粘贴到common.py文件的末尾
- #SimAM
- class SimAM(torch.nn.Module):
- def __init__(self, channels=None, out_channels=None, e_lambda=1e-4):
- super(SimAM, self).__init__()
- self.activaton = nn.Sigmoid()
- self.e_lambda = e_lambda
- def __repr__(self):
- s = self.__class__.__name__ + '('
- s += ('lambda=%f)' % self.e_lambda)
- return s
- @staticmethod
- def get_module_name():
- return "simam"
- def forward(self, x):
- b, c, h, w = x.size()
- n = w * h - 1
- x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)
- y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5
- return x * self.activaton(y)
如下图所示:

第②步:在yolo.py文件里的parse_model函数加入类名
首先找到yolo.py里面parse_model函数的这一行

然后把刚才加入的类SimAM添加到这个注册表里面:

第③步:创建自定义的yaml文件
首先在models文件夹下复制yolov5s.yaml 文件,粘贴并重命名为 yolov5s_SimAM.yaml

接着修改 yolov5s_SimAM.yaml ,将SimAM模块加到我们想添加的位置。
这里我先介绍第一种,第一种是将SimAM模块放在backbone部分的最末端,这样可以使注意力机制看到整个backbone部分的特征图,将具有全局视野,类似于一个小transformer结构。
将 [-1,1,SimAM,[1024]]添加到 SPPF 的下一层。即下图中所示位置:

同样的下面的head也得修改:

这里我们要把后面两个Concat的from系数分别由[ − 1 , 14 ] , [ − 1 , 10 ]改为[ − 1 , 15 ],[ − 1 , 11 ]。
然后将Detect原始的from系数[ 17 , 20 , 23 ]要改为[ 18 , 21 , 24 ] 。

yolov5s_SimAM.yaml 完整代码:
- # YOLOv5 🚀 by YOLOAir, GPL-3.0 license
- # Parameters
- nc: 80 # number of classes
- depth_multiple: 0.33 # model depth multiple
- width_multiple: 0.50 # layer channel multiple
- anchors:
- - [10,13, 16,30, 33,23] # P3/8
- - [30,61, 62,45, 59,119] # P4/16
- - [116,90, 156,198, 373,326] # P5/32
- backbone:
- # [from, number, module, args]
- [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
- [-1, 3, C3, [128]],
- [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
- [-1, 6, C3, [256]],
- [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
- [-1, 9, C3, [512]],
- [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
- [-1, 3, C3, [1024]],
- [-1, 1, SPPF, [1024, 5]], # 9
- [-1, 3, SimAM, [1024]], # 10
- ]
- # YOLOv5 v6.0 head
- head:
- [[-1, 1, Conv, [512, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 6], 1, Concat, [1]], # cat backbone P4
- [-1, 3, C3, [512, False]], # 14
- [-1, 1, Conv, [256, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 4], 1, Concat, [1]], # cat backbone P3
- [-1, 3, C3, [256, False]], # 18 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]],
- [[-1, 15], 1, Concat, [1]], # cat head P4
- [-1, 3, C3, [512, False]], # 21 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]],
- [[-1, 11], 1, Concat, [1]], # cat head P5
- [-1, 3, C3, [1024, False]], # 24 (P5/32-large)
- [[18, 21, 24], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
- ]
第④步:验证是否加入成功
在yolo.py 文件里面配置改为我们刚才自定义的yolov5s_SimAM.yaml


然后我们运行yolo.py

哇哦~果真是无参数捏!
第⑤步:修改train.py中 ‘--cfg’默认参数
我们先找到 train.py 文件的parse_opt函数,然后将第二行‘--cfg’的 default改为yolov5s_SimAM.yaml,然后就可以开始训练啦~

🚀三、在C3后添加SimAM注意力机制方法
第二种是将SimAM放在backbone部分每个C3模块的后面,这样可以使注意力机制看到局部的特征,每层进行一次注意力,可以分担学习压力。
步骤和方法1相同,只是yaml文件不同。
所以接下来只放修改yaml文件的部分~
第③步:创建自定义的yaml文件
将SimAM模块放在每个C3模块的后面,要注意通道的变化。
如下图所示:

同样的,下面的head部分也要做相应的修改
(数的时候一定要认真,我今天上午弄错了,卡了半天55~)

第二种方法的 yolov5s_SimAM.yaml 完整代码:
- # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
- # Parameters
- nc: 80 # number of classes
- depth_multiple: 0.33 # model depth multiple
- width_multiple: 0.50 # layer channel multiple
- anchors:
- - [10,13, 16,30, 33,23] # P3/8
- - [30,61, 62,45, 59,119] # P4/16
- - [116,90, 156,198, 373,326] # P5/32
- # YOLOv5 v6.0 backbone
- backbone:
- # [from, number, module, args]
- [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
- [-1, 3, C3, [128]],
- [-1, 3, SimAM, [128]],
- [-1, 1, Conv, [256, 3, 2]], # 4-P3/8
- [-1, 6, C3, [256]],
- [-1, 3, SimAM, [256]],
- [-1, 1, Conv, [512, 3, 2]], # 7-P4/16
- [-1, 9, C3, [512]],
- [-1, 3, SimAM, [512]],
- [-1, 1, Conv, [1024, 3, 2]], # 10-P5/32
- [-1, 3, C3, [1024]],
- [-1, 3, SimAM, [1024]],
- [-1, 1, SPPF, [1024, 5]], # 13
- ]
- # YOLOv5 v6.0 head
- head:
- [[-1, 1, Conv, [512, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 9], 1, Concat, [1]], # cat backbone P4
- [-1, 3, C3, [512, False]], # 17
- [-1, 1, Conv, [256, 1, 1]],
- [-1, 1, nn.Upsample, [None, 2, 'nearest']],
- [[-1, 6], 1, Concat, [1]], # cat backbone P3
- [-1, 3, C3, [256, False]], # 21 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]],
- [[-1, 18], 1, Concat, [1]], # cat head P4
- [-1, 3, C3, [512, False]], # 24 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]],
- [[-1, 14], 1, Concat, [1]], # cat head P5
- [-1, 3, C3, [1024, False]], # 27 (P5/32-large)
- [[21, 24, 27], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
- ]
第④步:验证是否加入成功
同样的方法,我们来运行一下yolo.py

OK~完成啦!
PS:今天尝试加入SimAM注意力机制训练时间为4个多小时(比昨天轻量化网络还快?);结果比之前尝试效果最佳的CA注意力机制还涨了近1个点,可以说非常奈斯了!