YOLOv5改进之添加注意力机制

本文主要给大家讲解一下&＃xff0c;如何在yolov5中添加注意力机制&＃xff0c;

这里提供SE通道注意力的改进方法&＃xff0c;其他注意力的添加方法&＃xff0c;大同小异

首先找到SE注意力机制的pytorch代码

class SELayer(nn.Module):
def __init__(self, c1, r&＃061;16):
super(SELayer, self).__init__()
self.avgpool &＃061; nn.AdaptiveAvgPool2d(1)
self.l1 &＃061; nn.Linear(c1, c1 // r, bias&＃061;False)
self.relu &＃061; nn.ReLU(inplace&＃061;True)
self.l2 &＃061; nn.Linear(c1 // r, c1, bias&＃061;False)
self.sig &＃061; nn.Sigmoid()
def forward(self, x):
b, c, _, _ &＃061; x.size()
y &＃061; self.avgpool(x).view(b, c)
y &＃061; self.l1(y)
y &＃061; self.relu(y)
y &＃061; self.l2(y)
y &＃061; self.sig(y)
y &＃061; y.view(b, c, 1, 1)
return x * y.expand_as(x)

直接把SE注意力机制的程序&＃xff0c;复制到models文件夹下面的common.py文件中

然后找到yolo.py&＃xff0c;在这个文件中找到下面这一行

 然后把SE添加到这个注册表里&＃xff0c;

 直接在C3Ghost后面加上SELayer

然后是要修改yaml文件&＃xff0c;对于SE这种即插即用的注意力机制&＃xff0c;输入和输出的通道数相同&＃xff0c;所以不会影响其他模块的运行

注意力机制也可以插在其他地方&＃xff0c;只要注意把通道数对应好就行

然后head部分也要进行相应的修改&＃xff0c;因为我们多加了一层&＃xff0c;所以需要保持head以及最后的输出层不变就得稍微改一下head部分&＃xff0c;同样&＃xff0c;注意力机制也可以放在head里面&＃xff0c;跟加在backbone里面的方法相同。

这是原始的head部分&＃xff0c;需要修改成下面这样

因为我们把SE注意力机制插在第八层之后&＃xff0c;所以原先的14层就变成了15层&＃xff0c;同样的&＃xff0c;最后的输出也需要把层数加一。

添加这些注意力机制是yolov5最基础的改进&＃xff0c;但是说实话&＃xff0c;注意力机制可以添加在很多的地方&＃xff0c;不一定会有效果&＃xff0c;所以插在哪里效果最好就需要大家自行讨论了。 注意力机制也有很多种&＃xff0c;se注意力机制应该是属于最基础的通道注意力了吧

来源：weixin_47060902