使用视觉注意力机制改进YOLOv5模型
1 前言
首先声明一下，本系列文章不是面向零基础的，在此之前需要会使用Pytorch、了解YOLOv5的结构、跑过YOLOv5的模型、使用过深度学习部署框架（例如OpenVINO、ONNX Runtime、TensorRT等），建议没有这些基础的同学，可以看上一个系列的教程。
相关的源码下载、环境安装，可以参照上一个系列文章的第一篇，这里建议使用Python3.8，因为很多Python3.6.5太老了。
上一个系列的教程是介绍YOLOv5怎么使用，可以不需要知道模型的结构，这个系列不行，必须知道 ，因为这个系列需要修改结构。
本套教程可以使用YOLOv5-6.0、6.1、6.2和7.0的代码。6.1我们上个系列使用了，而YOLOv5-6.2相对于6.1，增加了图像分类、自动编译量化、GPU基准性能（精度和速度）导出、单GPU复现、苹果设备（M1/M2）支持等功能，YOLOv5-7.0在此基础上，增加了实例分割、Comet 日志记录和可视化集成等功能，并且自称是世界上速度最快、精度最高的实例分割网络（截止2022年11月22日发布，击败了当时所有的SOTA模型）。当然，我们暂时用不到那么多花里胡哨的功能，我们只需要先学会怎么改检测网络就行。
2 模型结构配置文件解读
Ultralytics公司的YOLO系列模型都是根据配置文件搭建网络的，如果我们想对模型结构进行修改，那么必须知道配置文件和模型结构的对应关系。
打开models/yolov5s.yaml，此为yolov5s的配置文件，其中backbone和head内容如下：
# 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, 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
]
# 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]], # 13
[-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]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
本来自己想写一个配置文件解读的文章，但看到这篇文章——YOLOv5基础 | 深度解析yolov5s.yaml配置文件后，感觉写的十分完美，该文章是基于YOLOv5-6.0写的，配置文件的每个细节都有具体的介绍，对SPPF结构还做了相应的实验，可谓深入浅出，本人感觉无法做到比他更好了，非常推荐看这篇文章。
我觉得文章中对SPP层的解释还不够明朗。
我第一次看到SPP层的时候（当时还在学YOLOX），脑海中存在的疑惑：同一个输入张量，经过多个池化层之后，为何还能进行concat，难道三个池化核的输出特征图，高宽都一样吗？后面debug之后才知道，对于不同的池化核，可以通过调整padding和stride，使得三个池化核输出的特征图高宽一致，因此可以进行concat。对于SPPF结构，文章作者从实验角度演示了SPPF与SPP层等效，且运行时间只有后者的不到一半。
实验代码如下（为了简便运算，这里省略了池化前后的卷积层）：
import time
import torch
import torch.nn as nn
class SPP(nn.Module):
def __init__(self):
super().__init__()
self.maxpool1 = nn.MaxPool2d(5, 1, padding=2)
self.maxpool2 = nn.MaxPool2d(9, 1, padding=4)
self.maxpool3 = nn.MaxPool2d(13, 1, padding=6)
def forward(self, x):
o1 = self.maxpool1(x)
o2 = self.maxpool2(x)
o3 = self.maxpool3(x)
return torch.cat([x, o1, o2, o3], dim=1)
class SPPF(nn.Module):
def __init__(self):
super().__init__()
self.maxpool = nn.MaxPool2d(5, 1, padding=2)
def forward(self, x):
o1 = self.maxpool(x)
o2 = self.maxpool(o1)
o3 = self.maxpool(o2)
return torch.cat([x, o1, o2, o3], dim=1)
def main():
input_tensor = torch.rand(8, 32, 16, 16)
spp = SPP()
sppf = SPPF()
output1 = spp(input_tensor)
output2 = sppf(input_tensor)
print(torch.equal(output1, output2))
t_start = time.time()
for _ in range(100):
spp(input_tensor)
print(f"spp time: {time.time() - t_start}")
t_start = time.time()
for _ in range(100):
sppf(input_tensor)
print(f"sppf time: {time.time() - t_start}")
if __name__ == '__main__':
main()
我在本地也做了一次实验，结果如下：
True
spp time: 0.9689846038818359
sppf time: 0.333345890045166
可以看到SPPF层，运行时间只有SPP层的三分之一多一点。
3 模型组件代码解读
看完刚刚推荐的文章，我们对配置文件已经有了初步了解，接下来我们来看看模型各个组件的源码。
backbone和head中，共出现了Conv、C3、SPPF、nn.Upsample、Concat和Detect六种结构，其中Conv、C3、SPPF、Concat的代码都以类的形式保存在models/common.py中，nn.Upsample则本身就是类，Detect则保存在models/yolo.py中。当然，common.py包含的类远远不止上面几种，但限于篇幅，我们不可能全部都讲，我们这里只讲在标准YOLOv5s中出现的几种。
（1）Conv
这是YOLOv5最基本的零件，我们随便打开结构图的某一部分，都可以看见Conv类：
上面示意图中的ConvBNSiLU即为Conv，在模块名下面的k1, s1, p0, c64、k1, s1, p0, c64等，表示的是卷积的相关参数，例如，k1表示卷积核的大小为1，s1表示卷积核的步长为1，p0表示卷积时padding的像素个数，c64表示卷积核的输出通道为64。
从图中的模块名可以看到，它不但包含了卷积层，还包含了BN层和激活函数SiLU，我们找到Conv的代码：
class Conv(nn.Module):
# Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
default_act = nn.SiLU() # default activation
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self, x):
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
return self.act(self.conv(x))
初始化方法中的c1, c2分别是卷积层的输入、输出通道数；g和d分别是分组卷积与空洞卷积的参数，因为YOLOv5s的标准模型用不到，所以我们这里不展开，感兴趣的小伙伴自己查资料；act用来确定使用什么样的激活函数，它可以是布尔值，也可以是nn.Module的对象，如果act=True，则使用默认的激活函数，即SiLU，如果act=False，则使用nn.Identity()来代替，即没使用激活函数，如果act是nn.Module的对象，则说明这是自定义的激活函数。
代码中还出现了autopad函数，这个用来计算padding的该设置成多少，其代码如下：
def autopad(k, p=None, d=1): # kernel, padding, dilation
# Pad to 'same' shape outputs
if d > 1:
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
（2）C3
C3有两种，它们之间的区别在于所包含的BottleNeck是否有多连接（后面会具体介绍），两种C3的结构如下：
代码如下：
class C3(nn.Module):
# CSP Bottleneck with 3 convolutions
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
# self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
def forward(self, x):
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
初始化参数中，n表示Bottleneck模块的数量，shortcut参数表示是否进行短连接，g是分组卷积参数，e是通道扩展倍数。self.cv3的输入通道之所以是2 * c_，是因为它的输入是两个分支Concat的结果，这两个分支的输出均为c_。
（3）Bottleneck
C3的示意图和代码中还出现了Bottleneck，并且有两种，在backbone中使用的是Bottleneck1，在head中使用的是Bottleneck2，两者区别在于有没有短连接，它们的示意图如下：
类定义如下：
class Bottleneck(nn.Module):
# Standard bottleneck
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c2, 3, 1, g=g)
self.add = shortcut and c1 == c2
def forward(self, x):
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
初始化方法的参数含义和C3模块一致。
（4）SPPF
这个我们在前面演示过，原理已经懂了，但前面的程序不完整，这里给出完整的代码：
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
super().__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
（5）Concat
其实这个结构没必要讲，就是将torch.cat函数封装了一下，代码如下：
class Concat(nn.Module):
# Concatenate a list of tensors along dimension
def __init__(self, dimension=1):
super().__init__()
self.d = dimension
def forward(self, x):
return torch.cat(x, self.d)
（6）Detect
这个类不在models/common.py中，而是在models/yolo.py中，这个类写了很长，我们不可能把完整代码贴过来，我这边只把初始化方法贴过来。
class Detect(nn.Module):
# YOLOv5 Detect head for detection models
stride = None # strides computed during build
dynamic = False # force grid reconstruction
export = False # export mode
def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer
super().__init__()
self.nc = nc # number of classes
self.no = nc + 5 # number of outputs per anchor
self.nl = len(anchors) # number of detection layers
self.na = len(anchors[0]) // 2 # number of anchors
self.grid = [torch.empty(0) for _ in range(self.nl)] # init grid
self.anchor_grid = [torch.empty(0) for _ in range(self.nl)] # init anchor grid
self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2)
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
self.inplace = inplace # use inplace ops (e.g. slice assignment)
def forward(self, x):
pass
def _make_grid(self, nx=20, ny=20, i=0):
pass
初始化方法的参数中，nc是图像中目标的类别数，anchors是先验框的尺寸（就是配置文件里的anchors），ch是由每个检测头的输入通道数构成的列表。
在配置文件中，Detect只出现了一次，但实际上检测头有三个，这个类实际上定义了多个检测头，每个检测头都是self.m中的一个元素，我们可以debug一下看看：
In[2]: len(self.m)
Out[2]: 3
从YOLOv5-6.x开始，把原本属于后处理的anchorbox解析与高宽缩放移到模型内部处理了，就在Detect的forward方法中，这个说起来比较复杂了，我们这里不展开，知道这么回事就行。
（7）对模型对象进行Debug代码回溯
Detect类在models/yolo.py中，而不是在common.py中，我是怎么知道的，难道把所有源码都看一遍吗？项目这么大，我不可能把源码中的每个脚本都看一遍，深度学习算法工程师最忌讳的事情就是遍历源码，看源码是非常低效的事情，能不看源码尽量不看源码，比较推荐使用debug的方式，在关键节点打断点，通过断点栈帧回溯功能，找到源码中模型的创建位置。
我是在SPPF的__init__方法中打了断点，因为模型在创建过程中，肯定会执行SPPF的__init__方法。然后把数据配置好（参考上一个系列的第二篇文章），之后debug运行train.py（之所以运行train.py而不是detect.py，是因为我试了，debug运行detect.py时，SPPF的初始化方法不会中断，detect不行之后我才尝试train的），程序中断后，根据断点栈帧的Frames窗口，看是谁在调用它，结果发现是models/yolo.py的parse_model(d, ch)方法在调用它。
接着把SPPF__init__方法中的断点去掉，在parse_model(d, ch)中调用SPPF__init__方法的语句上打断点：
然后点击不断点击Resume Program（F9）按钮，直到变量窗口中Special Variables下的m变成<class '__main__.Detect'>为止：
点击截图中的3，进入到Detect的初始化方法，我们就能看见Detect类在哪里了：
除此之外，我们根据断点Frames，也可以看到创建模型是在train.py的哪一条语句中
使用Pycharm的Debug Frames，可以知道程序的调用关系，掌握这种回溯功能，使得我们分析程序，不需要从头开始一步步Debug，只需要找到关键节点，在关键节点上打断点就可以。
比如，我想知道train.py是在哪一步构建了模型，假如我事先不知道模型类Model在哪里，或者说我压根就不知道模型类的名称是啥，是Model，还是Net，还是Yolov5，又或者是其他名字，我总不可能把整个项目的代码都看一遍吧？但我知道新版的YOLOv5s它肯定使用了SPPF，那我就在SPPF的初始化方法上打断点，然后debug执行train.py程序，根据断点栈帧（Frames），反向确定YOLOv5s的模型是在哪个地方创建的，使用的类名是什么，类定义在哪里等。
我个人认为，源码是无底洞，我之前尝试过去看YOLOX的全工程源码，但源码经过层层封装，到底层就看不懂程序在干什么了。其实，任何一个大型软件工程项目，背后都是有一个团队在开发，团队的每个成员分工合作才有了这个项目。我们作为学习者，很难凭一人之力去对抗整个团队，一条一条地看源码就是在用有限的生命去追逐无限的细节。
那我们该如何学习新模型呢？学习一个新的模型，最应该做的事就是去看论文或者教程（博客或者视频，站在巨人的肩膀上），知道它的大致原理，然后看项目下面的readme，搞懂这个东西该怎么用，API是什么，那么大致就OK了。
如果想要改进别人的模型，该怎么弄呢？首先，需要明确的是，我们要改的模型必须是经典模型，比如YOLO系列，有改的价值，并且经典模型的代码工程化做的比较好，网上有比较多代码的解读，容易找到配置文件（或者其他容易入手的文件，可以称之为切入点文件），改起来也比较容易。对于非经典模型，特别是很难找到相关文档教程的模型，你想搞清楚它的结构（可能你都不知道这个模型是否使用了配置文件），那么只能从train.py一行一行读代码了，这个过程将是个噩梦。
对于源码，当我们想要了解某个具体细节（比如NMS）的时候，才推荐使用debug进入到源码的相应部分去看，这也不是全看源码，而是debug源码的其中一部分。
4 模块集成
经过上节的讲解，我们已经对配置文件有了基本的了解，现在我们来看看如何根据配置文件中的信息创建各个模块。
把配置文件中的各个模块串起来，是在models/yolo.py的parse_model(d, ch)方法中进行的，我们在该方法的第二条语句打断点：
随后使用debug运行train.py
parse_model(d, ch)方法中的ch参数是输入通道（通过断点栈帧回溯到Model类的初始化方法，注释中有说这个参数是输入通道数），而d则是从配置文件中解析出来的配置信息。进入parse_model函数后，先是从配置信息中获取anchor的个数，模型的输出通道（即每个预测框包含的变量个数），代码如下：
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
接下来是通过循环解析backbone和head的参数，内容如下：
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
m = eval(m) if isinstance(m, str) else m # eval strings
接下来我们debug：
进入循环后，每次从配置信息中获取一条参数（f, n, m, args）。
当执行完m = eval(m) if isinstance(m, str) else m后，m将变成类对象：
接下来的for循环是将args中的内容（有可能为字符串）转换成python中的数据类型（或者对象）。
for j, a in enumerate(args):
try:
args[j] = eval(a) if isinstance(a, str) else a # eval strings
except NameError:
pass
如果m是C3或者Conv，那么需要在args前面加上输入通道和输出通道，如果是其他模块，则根据相应的规则组织参数。
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
c1, c2 = ch[f], args[0] # 当i为0时，及解析的为backbone的首行，则f为-1，ch为[3]，args[0]为64
if c2 != no: # if not output
c2 = make_divisible(c2 * gw, 8)
args = [c1, c2, *args[1:]]
if m in [BottleneckCSP, C3, C3TR, C3Ghost]:
args.insert(2, n) # number of repeats
n = 1
elif m is nn.BatchNorm2d:
args = [ch[f]]
elif m is Concat:
c2 = sum(ch[x] for x in f)
elif m is Detect:
args.append([ch[x] for x in f])
if isinstance(args[1], int): # number of anchors
args[1] = [list(range(args[1] * 2))] * len(f)
elif m is Contract:
c2 = ch[f] * args[0] ** 2
elif m is Expand:
c2 = ch[f] // args[0] ** 2
else:
c2 = ch[f]
总之，在m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)执行前，args已经按照各个模块的初始化方法配置好参数了。
接下来是创建模块，并把每个模块的通道数存入到指定列表：
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
t = str(m)[8:-2].replace('__main__.', '') # module type
np = sum(x.numel() for x in m_.parameters()) # number params
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}') # print
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
layers.append(m_)
if i == 0:
ch = []
ch.append(c2)
这里m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)是创建模块；
i是模块序号，f是数据来源，type是模块类型，np是参数总量，m_.i, m_.f, m_.type, m_.np = i, f, t, np是将这些变量存到模型中；
随后save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)是将数据来源f保存到列表save中，这里f也有可能是列表，例如[[-1, 6], 1, Concat, [1]]，x % i是为了保证x不大于i，也就是f或f中的元素不会大于i，若f为-1，则不存入到save中；
layers.append(m_)是将模块保存；
ch.append(c2)则是保存当前模块的输出通道，它是后续模块的输入通道，后续模块可以通过ch[f]查询得到。
最后函数返回两个值，一是将layers整合成nn.Sequential的结果，二是save中的序号排序后的结果。
return nn.Sequential(*layers), sorted(save)
这里save的作用是什么，为什么要对其进行排序，单单看parse_model(d, ch)函数是看不出来的，这个要讲明白，需要分析整个Model类，个人感觉没必要。
本节介绍了配置信息解析函数parse_model(d, ch)，该函数根据配置信息创建各个模块，但并没有把这些模块串起来，这个需要分析Model类的forward方法，这个工作量太大，且对我们修改网络结构帮助不大，我们只需要大致知道配置信息进入到for循环后是怎么处理的，走的是if的哪个分支，又是怎么根据处理后的参数创建模块的，知道这些足够了。
5 总结
本文是这个系列文章中最难最琐碎的，搞明白了如何根据配置文件创建YOLOv5s模型的各个模块，之后添加注意力机制就会变得很容易。
扫一扫
4万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
