YOLOv5改进之添加注意力机制

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

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

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

class SELayer(nn.Module):
    def __init__(self, c1, r=16):
        super(SELayer, self).__init__()
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.l1 = nn.Linear(c1, c1 // r, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.l2 = nn.Linear(c1 // r, c1, bias=False)
        self.sig = nn.Sigmoid()

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avgpool(x).view(b, c)
        y = self.l1(y)
        y = self.relu(y)
        y = self.l2(y)
        y = self.sig(y)
        y = y.view(b, c, 1, 1)
        return x * y.expand_as(x)

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

然后找到yolo.py，在这个文件中找到下面这一行