YOLO算法改进5【中阶改进篇】：添加SENet注意力机制

SE-Net 是 ImageNet 2017（ImageNet 收官赛）的冠军模型，是由WMW团队发布。具有复杂度低，参数少和计算量小的优点。且SENet 思路很简单，很容易扩展到已有网络结构如 Inception 和 ResNet 中。
我们可以看到，已经有很多工作在空间维度上来提升网络的性能。那么很自然想到，网络是否可以从其他层面来考虑去提升性能，比如考虑特征通道之间的关系？作者基于这一点并提出了Squeeze-and-Excitation Networks（简称SE-Net）。在该结构中，Squeeze和Excitation是两个非常关键的操作，所以以此来命名。作者出发点是希望建立特征通道之间的相互依赖关系。并未引入一个新的空间维度来进行特征通道间的融合，而是采用了一种全新的“特征重标定”策略。具体来说，就是通过学习的方式来自动获取到每个特征通道的重要程度，然后依照这个重要程度去提升有用的特征并抑制对当前任务用处不大的特征。

一、不改变原网络深度的改进方法

首先是打开models/yolov5s.yaml文件，我们在backbone中的SPPF之前增加SENet。增添位置如下，是将backbone中第4个C3模块替换为SE_Block，如上图。需要注意的是通道数要匹配，SENet并不改变通道数，由于原C3的输出通道数为1024*0.5=512，所以我们这里的写的是1024，这里的1024是传入到上面我们定义的Class SE_Block(nn.Moudel)中的c2参数，c1参数是由上一层的输出通道数控制的。参考链接

1.添加SENet.yaml文件
添加至/models/文件中

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2  conv1(3,32,k=6,s=2,p=2)
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4  conv2(32,64,k=3,s=2,p=1)
   [-1, 3, C3, [128]],  # C3_1 有Bottleneck
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8  conv3(64,128,k=3,s=2,p=1)
   [-1, 6, C3, [256]], # C3_2 Bottleneck重复两次
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16  conv4(128,256,k=3,s=2,p=1)
   [-1, 9, C3, [512]], # C3_3 Bottleneck重复三次 输出256通道
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32   Conv5(256,512,k=3,s=2,p=1)
   #[-1, 3, C3, [1024]],  # C3_4 Bottleneck重复1次  输出512通道
   [-1, 1, SE_Block, [1024]],  # 增加通道注意力机制 输出为512通道
   [-1, 1, SPPF, [1024, 5]],  # 9  每个都是K为5的池化
  ]

# YOLOAir 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, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

2.common配置
在models/common.py文件中增加以下代码

• 上图是作者提出的SE模块的示意图。给定一个输入$x$，其特征通道数为$c_1$，通过一系列卷积变换后得到一个特征通道数为$c_2$的特征。与传统的CNN不一样的是，接下来将通过三个操作来重标定前面得到的特征。
• 首先是Squeeze操作，顺着空间维度来进行特征压缩，将每个二维的特征通道变成一个实数，这个实数某种程度上具有全局的感受野，并且输出的维度和输入的特征通道数相匹配。它表征着在特征通道上响应的全局分布，而且使得靠近输入的层也可以获得全局的感受野，这一点在很多任务中都是非常有用的。
• 其次是Excitation操作，它是一个类似于循环神经网络中门的机制。通过参数来为每个特征通道生成权重，其中参数被学习用来显式地建模特征通道间的相关性。
• 最后是一个Reweight的操作，我们将Excitation的输出的权重看做是进过特征选择后的每个特征通道的重要性，然后通过乘法逐通道加权到先前的特征上，完成在通道维度上的对原始特征的重标定。
  ——————————————————————————————————————————
  
• 这里的注意力机制想法非常简单，即针对每一个 channel 进行池化处理，就得到了 channel
  个元素，通过两个全连接层，得到输出的这个向量。值得注意的是，第一个全连接层的节点个数等于 channel 个数的 $\frac{1}{4}$（论文作者发现如果将第一个全连接层的节点个数替换成原来的$\frac{1}{4}$，可以在参数数量适度增加的情况下提高准确性，而且并没有明显的延迟。），然后第二个全连接层的节点就和channel 保持一致。这个得到的输出就相当于对原始的特征矩阵的每个 channel 分析其重要程度，越重要的赋予越大的权重，越不重要的就赋予越小的权重。
• 就拿上图来说，首先对四个通道进行平均池化得到四个值，然后经过两个全连接层之后得到通道权重的输出。等权重输出以后，则将对应通道的权重乘以原来的特征矩阵就得到了新的特征矩阵，以上便是SE模块的详细实现过程。

class SE_Block(nn.Module):
    def __init__(self, c1, c2):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 平均池化
        self.fc = nn.Sequential(
            nn.Linear(c1, c2 // 16, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(c2 // 16, c2, bias=False),
            nn.Sigmoid()
        )
 
    def forward(self, x):
        # 添加注意力模块
        b, c, _, _ = x.size()  # 分别获取batch_size,channel
        y = self.avg_pool(x).view(b, c)  # y的shape为【batch_size, channels】
        y = self.fc(y).view(b, c, 1, 1)  # shape为【batch_size, channels, 1, 1】
        out = x * y.expand_as(x)  # shape 为【batch, channels,feature_w, feature_h】
        return out
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

3.yolo.py配置
找到 models/yolo.py 文件中 parse_model() 类，在列表中添加SE_Block，这样可以获得我们要传入的参数。

if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
                 BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, SE_Block]:
1
2

二、改变原网络深度的改进方法

比如我要在第一个C3后面加一个SE。yaml的修改如下。接下来稍微麻烦一点了【需要你了解v5的每层结构】，由于我们在backbone中加入了一层，也就是相当于后面的网络与之前相比都往后移动了一层，那么在后面的Concat部分中融合的特征层的索引也会收到影响，因此我们需要的是修改Concat层的from参数。参考链接

1.添加SENet.yaml文件
添加至/models/文件中

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2  conv1(3,32,k=6,s=2,p=2)
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4  conv2(32,64,k=3,s=2,p=1)
   [-1, 3, C3, [128]],  # C3_1 有Bottleneck
   [-1, 1, SE_Block, [128]],  # 增加通道注意力机制 输出为512通道
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8  conv3(64,128,k=3,s=2,p=1)
   [-1, 6, C3, [256]], # C3_2 Bottleneck重复两次
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16  conv4(128,256,k=3,s=2,p=1)
   [-1, 9, C3, [512]], # C3_3 Bottleneck重复三次 输出256通道
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32   Conv5(256,512,k=3,s=2,p=1)
   [-1, 3, C3, [1024]],  # C3_4 Bottleneck重复1次  输出512通道
   [-1, 1, SPPF, [1024, 5]],  # 9  每个都是K为5的池化
  ]
"""可以看到实际就是每个Concat也后面移动一层，因此yaml修改为一下。最终的Detect的from也需要修改。""
# YOLOAir v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],  # conv1(512,256,1,1)
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 7], 1, Concat, [1]],  # cat backbone P4  将C3_3与SPPF出来后的上采样拼接 拼接后的通道为512
   [-1, 3, C3, [512, False]],  # 13  conv(256,256,k=1,s=1)  没有残差边
 
   [-1, 1, Conv, [256, 1, 1]], # conv2(256,128,1,1)
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 5], 1, Concat, [1]],  # cat backbone P3  与C3_2拼接，输出256通道
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small) conv3(128,128,1,1)
 
   [-1, 1, Conv, [256, 3, 2]],# conv4(128,128,3,2,1)
   [[-1, 15], 1, Concat, [1]],  # cat head P4  拼接后256通道
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)  conv5(256,256,1,1)
 
   [-1, 1, Conv, [512, 3, 2]],# conv6(256,256,3,2,1)
   [[-1, 11], 1, Concat, [1]],  # cat head P5  拼接后是512
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
 
   [[18, 21, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

2.common配置
在models/common.py文件中增加以下代码

class SE_Block(nn.Module):
    def __init__(self, c1, c2):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 平均池化
        self.fc = nn.Sequential(
            nn.Linear(c1, c2 // 16, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(c2 // 16, c2, bias=False),
            nn.Sigmoid()
        )
 
    def forward(self, x):
        # 添加注意力模块
        b, c, _, _ = x.size()  # 分别获取batch_size,channel
        y = self.avg_pool(x).view(b, c)  # y的shape为【batch_size, channels】
        y = self.fc(y).view(b, c, 1, 1)  # shape为【batch_size, channels, 1, 1】
        out = x * y.expand_as(x)  # shape 为【batch, channels,feature_w, feature_h】
        return out
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

3.yolo.py配置
找到 models/yolo.py 文件中 parse_model() 类，在列表中添加SE_Block，这样可以获得我们要传入的参数。

if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
                 BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, SE_Block]:
1
2

4.训练模型

python train.py --cfg SENet.yaml

1
2