写在开头:
最近学习了魔傀面具老师的YOLO系列课程,以免自己忘记一些步骤,特地在CSDN里写下这篇博文,权当作是一篇笔记,温故知新。
魔傀面具老师的b站链接:https://space.bilibili.com/286900343?spm_id_from=333.337.0.0
github:https://github.com/z1069614715
一、修改的文件及其解释
1、yolov5s.yaml:构建模型所需的参数列表。
图1. yolov5s.yaml文件中参数的解释
图1中,各个参数的含义已在图上标注。值得注意的是,backbone处的参数解释,与head处的参数保持一致。
重点解释backbone的参数,以第一个list为例:
[-1, 1, Conv, [64, 6, 2, 2]]
-1:与前一层相连,上一层的输出作为当前层的输入。
1:当前层的重复次数,在yolo.py文件中,以变量n表示。
Conv:当前层所属模块的类型,在yolo.py文件中,以变量m表示。包含的类型如下所示:
if m in {Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
[64, 6, 2, 2]:当前层所属模块的参数,在yolo.py文件中,被存储在args变量中。
2、yolo.py:构建模型的类和各个函数。
以parse_model函数为主,源码标注了每个变量的英文,比较容易看懂。单步调试parse_model函数,结合注释能更好地了解每个变量的变化,从而理解构建模型的步骤。
此外,需要强调的是args变量,args以list格式存储了yolov5s.yaml文件的模块变量参数,如:[64, 6, 2, 2]。
图2. args变量变化。左:变化前。右:变化后
然而,此时只是初始化参数,并不能满足模块的参数需要,以当前层所属模块Conv为例:
图3. Conv模块源码
Conv模块至少需要输入c1、c2、k、s、p这5个参数,参数含义如下:
c1:channel_in,输入通道数
c2:channel_out,输出通道数
k:kernel size,卷积核大小
s:stride,卷积步距
p:padding,卷积填充数
初始输入通道数为3,传入c1变量,将初始args变量的第一个参数传入c2变量。
args = [c1, c2, *args[1:]]
如图2所示,将输入通道数、输出通道数、卷积核大小、卷积布局、卷积填充数写入args变量中。
图4. 第一层Conv参数
以第一层Conv为例,其余模块原理类似。
3、common.py:各个模块的构造函数。
有如下模块:
Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x
二、添加注意力机制的步骤
1、无需修改通道数的注意力
以SimAM注意力机制为例,源码:https://github.com/z1069614715/objectdetection_script/blob/master/cv-attention/SimAM.py
1.1、在models文件下新建脚本SimAM.py,写入源码。
出于个人习惯,在models文件夹下又新建了一个Attention文件夹,并将注意力机制源码脚本写在子文件夹中。
1.2、在yolo.py文件中,导入SimAM脚本。
图5. 导入注意力机制
1.3、在yolov5s.yaml文件中,在需要添加注意力机制的地方,添加代码,主要修改module和args参数:
[-1, -1, SimAM, [1e-4]],
其中,1e-4:SimAM函数的唯一参数。
1.4、修改各层所连接的层数
当添加了注意力机制后,每层的序号就发生了变化,此时需相应地改变yolov5s.yaml文件中,各个list的连接层数。
在各个Concat层中,当前Concat层需要连接上一层和指定层,分别用-1和指定层序号表示。添加注意力机制后,-1保持不变,指定层序号改变。
例如:只添加一层注意力机制,则在cat head P4层,变化为:[-1, 14] -> [-1, 15]。再比如,在Detect层,变化为:[17, 20, 23] -> [18, 21, 24]。
1.5、运行,验证
图6. 无需通道数注意力机制的添加结果
2、需要修改通道数的注意力
以SEAttention为例,源码:https://github.com/z1069614715/objectdetection_script/blob/master/cv-attention/SE.py
2.1、大致步骤与“无需修改通道数的注意力”相同,除了传入通道数,需要额外添加一行代码:
elif m in [SEAttention]: args = [ch[f], *args]
[SEAttention]可扩展为:[SEAttention, ShuffleAttention, ......]
2.2、更改yolov5s.yaml文件参数:
其中,参数[16]:SEAttention函数中,除了通道数外的唯一参数。
2.3、运行,验证:
图7. 需输入通道数的注意力机制的添加结果
三、总结
总体步骤可分为三步:
- 构造并导入注意力机制脚本
- 修改模型参数yaml文件
- 在构建模型yolo.py文件中,传入通道数至注意力层中。
总的来说,步骤比较简单,目前尝试的注意力机制有SimAM、ShuffleAttention、SEAttention三种,后续会尝试更多注意力机制,如:BAM、CBAM等。