YOLOv5 加入SE注意力机制

步骤1：在common.py中添加SE模块

class SE(nn.Module):
    def __init__(self, c1, c2, r=16):
        super(SE, 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):
        print(x.size())
        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)

 步骤2：在yolo.py文件中加入类名

 新建yolos_se.yaml

注意力机制可以加在Backbone、Neck、Head等部分，常见的有两种：一种是在主干的SPPF前面添加一层；二是将Backbone中的C3全部替换。 

yolo.py

 

 运行yolo.py

 修改train.py中的'--cfg'默认参数