YOLO算法改进指南：8.添加SimAM注意力机制 ｜无参数注意力模块

本文是中山大学在注意力机制方面的尝试，从神经科学理论出发，构建了一种能量函数挖掘神经元重要性，并对此推导出了解析解以加速计算。通过ImageNet分类、COCO检测与分割等任务验证了所提SimAM的灵活性与有效性。值得一提的是，所提SimAM是一种无参数注意力模块。

• 代码地址：https://github.com/ZjjConan/SimAM

一、SimAM详解

1. 概述

本文提出一种概念简单且非常有效的注意力模块。不同于现有的通道/空域注意力模块，该模块无需额外参数为特征图推导出3D注意力权值。具体来说，本文基于著名的神经科学理论提出优化能量函数以挖掘神经元的重要性。本文进一步针对该能量函数推导出一种快速解析解并表明：该解析解仅需不超过10行代码即可实现。该模块的另一个优势在于：大部分操作均基于所定义的能量函数选择，避免了过多的结构调整。最后，本文在不同的任务上对所提注意力模块的有效性、灵活性进行验证。

本文主要贡献包含以下几点：

• 受启发于人脑注意力机制，本文提出一种3D注意力模块并设计了一种能量函数用于计算注意力权值；
• 本文推导出了能量函数的解析解加速了注意力权值的计算并得到了一种轻量型注意力模块；
• 将所提注意力嵌入到现有ConvNet中在不同任务上进行了灵活性与有效性的验证。

2. 理论

在正式介绍本文所提注意力模块之前，我们先对现有代表性注意力模块(比如SE、CBAM、GC)进行简要总结；然后，我们再引出本文所提完全不同架构的注意力模块。

SimAM注意力机制原理图：

上图a与b列出了现有两种类型的注意力模块：

• 通道注意力：1D注意力，它对不同通道区别对待，对所有位置同等对待；
• 空域注意力：2D注意力，它对不同位置区别对待，对所有通道同等对待。

以下图为例，SE缺失了关于"grey_whale"的某些重要成分。我们认为3D注意力比1D和2D更佳，进而提出了上图c的3D注意力模块。

现有注意力模块的另一个重要影响因素：权值生成方法。现有注意力往往采用额外的子网络生成注意力权值，比如SE的GAP+FC+ReLU+FC+Sigmoid。更多注意力模块的操作、参数量可参考下表。总而言之，现有注意力的结构设计需要大量的工程性实验。我们认为：注意力机制的实现应当通过神经科学中的某些统一原则引导设计。

📌 本文亮点总结：

1. 现有注意力模块的另一个重要影响因素：权值生成方法。现有注意力往往采用额外的子网络生成注意力权值，比如SE的GAP+FC+ReLU+FC+Sigmoid。
2. 在神经科学中，信息丰富的神经元通常表现出与周围神经元不同的放电模式。而且，激活神经元通常会抑制周围神经元，即空域抑制。换句话说，具有空域抑制效应的神经元应当赋予更高的重要性。最简单的寻找重要神经元的方法：度量神经元之间的线性可分性。

二、YOLOv5改进

1. 增加SimAM.yaml文件

# 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
   [-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
  ]

# 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)
   [-1, 1, SimAM, [1024]], #修改

   [[17, 20, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

2. common.py配置

在./models/common.py文件中增加以下内容

import torch
import torch.nn as nn

class SimAM(torch.nn.Module):
    def __init__(self, channels=None, e_lambda=1e-4):
        super(SimAM, self).__init__()

        self.activaton = nn.Sigmoid()  # Sigmoid 激活函数的实例
        self.e_lambda = e_lambda  # 正则化项的参数 lambda

    def __repr__(self):
        s = self.__class__.__name__ + '('
        s += ('lambda=%f)' % self.e_lambda)  # 返回表示模型的字符串，包括 lambda 参数的值
        return s

    @staticmethod
    def get_module_name():
        return "simam"  # 返回模块的名称，为 "simam"

    def forward(self, x):
        # 前向传播函数，对输入 x 进行处理并返回结果

        b, c, h, w = x.size()    # 获取输入张量的形状信息
        n          = w * h - 1   # 计算总像素数减1

        x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)  # 计算每个像素与均值的差的平方
        y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5
        # 计算激活函数的输入 y，使用 SimAM 公式：x_minus_mu_square / (4 * (均值方差 + 正则化项)) + 0.5

        return x * self.activaton(y)  # 返回经过激活函数后的结果与输入张量 x 的乘积

3. yolo.py配置

找到models/yolo.py文件中parse_model()函数的for i, (f, n, m, args) in enumerate(d['backbone'] + d['head'])（258行上下）并其循环内添加如下代码。

elif m is SimAM:
    c1, c2 = ch[f], args[0]
    if c2 != no:
        c2 = make_divisible(c2 * gw, 8)

4. 训练模型

python train.py --cfg yolov5_SimAM.yaml