YOLOX 改进 004：Backbone 添加 SimAM 注意力机制

论文题目：《SimAM: A Simple, Parameter-Free Attention Module for Convolutional Neural Networks》

论文地址：https://proceedings.mlr.press/v139/yang21o/yang21o.pdf

官方源码：GitHub - ZjjConan/SimAM

YOLOX 官方源码：GitHub - Megvii-BaseDetection/YOLOX

YOLOX 网络解读参考：CSDN - Pytorch 搭建 YOLOX 目标检测平台

1. SimAM 简介

SimAM（Simple Attention Module）是一种轻量级的注意力机制，旨在增强神经网络的表示能力，而不会显著增加计算开销。它通过对特征图进行自适应加权，来提高网络对关键特征的关注度。

SimAM 的基本思想是使用一个简单的模块来捕捉特征图中的重要信息。与其他复杂的注意力机制（如 SE 模块、CBAM 等）相比，SimAM 通过引入一个自适应加权机制，实现了高效的特征增强。

SimAM 的优点：

•  计算开销低： SimAM 不引入额外的卷积层或全连接层，计算量相对较低。
• 易于实现： SimAM 的实现相对简单，可以很容易地嵌入到现有的卷积神经网络中。
• 有效性： 实验表明，SimAM 在提高网络性能方面表现出色，能够显著提高图像分类、目标检测等任务的准确性。 

2. 项目环境

• 解释器：3.9.19
• 框架：Pytorch 2.0.0 + CUDA 11.8
• 系统：Win10 / Ubuntu 20.04

3. 核心代码

import torch
import torch.nn as nn

__all__ = ['SimAM']

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

        self.activaton = nn.Sigmoid()
        self.e_lambda = e_lambda

    def __repr__(self):
        s = self.__class__.__name__ + '('
        s += ('lambda=%f)' % self.e_lambda)
        return s

    @staticmethod
    def get_module_name():
        return "simam"

    def forward(self, x):
        b, c, h, w = x.size()

        n = w * h - 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

        return x * self.activaton(y)

4. 添加方法  

 第 1 步 ：在 YOLOX-main/yolox/models/ 目录下新建目录 add_modules/，在 YOLOX-main/yolox/models/add_modules/ 目录下新建 Python 源文件 SimAM.py，将以上 SimAM 注意力模块的核心代码复制粘贴至 SimAM.py 文件中。

 第 2 步 ：定位到 YOLOX-main/yolox/models/ 目录下的 darknet.py 文件，该文件中的内容为 YOLOX 主干网络的源代码。

以下代码为我在主干网络（Backbone）添加完 SimAM 注意力机制之后的最终代码，大家需要参照代码中的注释一步一步地修改 darknet.py 文件中的代码。

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.

from torch import nn

from .network_blocks import BaseConv, CSPLayer, DWConv, Focus, ResLayer, SPPBottleneck

from add_modules.SimAM import SimAM

class Darknet(nn.Module):
    # number of blocks from dark2 to dark5.
    depth2blocks = {21: [1, 2, 2, 1], 53: [2, 8, 8, 4]}

    def __init__(
        self,
        depth,
        in_channels=3,
        stem_out_channels=32,
        out_features=("dark3", "dark4", "dark5"),
    ):
        """
        Args:
            depth (int): depth of darknet used in model, usually use [21, 53] for this param.
            in_channels (int): number of input channels, for example, use 3 for RGB image.
            stem_out_channels (int): number of output chanels of darknet stem.
                It decides channels of darknet layer2 to layer5.
            out_features (Tuple[str]): desired output layer name.
        """
        super().__init__()
        assert out_features, "please provide output features of Darknet"
        self.out_features = out_features
        self.stem = nn.Sequential(
            BaseConv(in_channels, stem_out_channels, ksize=3, stride=1, act="lrelu"),
            *self.make_group_layer(stem_out_channels, num_blocks=1, stride=2),
        )
        in_channels = stem_out_channels * 2  # 64

        num_blocks = Darknet.depth2blocks[depth]
        # create darknet with `stem_out_channels` and `num_blocks` layers.
        # to make model structure more clear, we don't use `for` statement in python.
        self.dark2 = nn.Sequential(
            *self.make_group_layer(in_channels, num_blocks[0], stride=2)
        )
        in_channels *= 2  # 128
        self.dark3 = nn.Sequential(
            *self.make_group_layer(in_channels, num_blocks[1], stride=2)
        )
        in_channels *= 2  # 256
        self.dark4 = nn.Sequential(
            *self.make_group_layer(in_channels, num_blocks[2], stride=2)
        )
        in_channels *= 2  # 512

        self.dark5 = nn.Sequential(
            *self.make_group_layer(in_channels, num_blocks[3], stride=2),
            *self.make_spp_block([in_channels, in_channels * 2], in_channels * 2),
        )

    def make_group_layer(self, in_channels: int, num_blocks: int, stride: int = 1):
        "starts with conv layer then has `num_blocks` `ResLayer`"
        return [
            BaseConv(in_channels, in_channels * 2, ksize=3, stride=stride, act="lrelu"),
            *[(ResLayer(in_channels * 2)) for _ in range(num_blocks)],
        ]

    def make_spp_block(self, filters_list, in_filters):
        m = nn.Sequential(
            *[
                BaseConv(in_filters, filters_list[0], 1, stride=1, act="lrelu"),
                BaseConv(filters_list[0], filters_list[1], 3, stride=1, act="lrelu"),
                SPPBottleneck(
                    in_channels=filters_list[1],
                    out_channels=filters_list[0],
                    activation="lrelu",
                ),
                BaseConv(filters_list[0], filters_list[1], 3, stride=1, act="lrelu"),
                BaseConv(filters_list[1], filters_list[0], 1, stride=1, act="lrelu"),
            ]
        )
        return m

    def forward(self, x):
        outputs = {}
        x = self.stem(x)
        outputs["stem"] = x
        x = self.dark2(x)
        outputs["dark2"] = x
        x = self.dark3(x)
        outputs["dark3"] = x
        x = self.dark4(x)
        outputs["dark4"] = x
        x = self.dark5(x)
        outputs["dark5"] = x
        return {k: v for k, v in outputs.items() if k in self.out_features}


class CSPDarknet(nn.Module):
    def __init__(
        self,
        dep_mul,
        wid_mul,
        out_features=("dark3", "dark4", "dark5"),
        depthwise=False,
        act="silu",
    ):
        super().__init__()
        assert out_features, "please provide output features of Darknet"
        self.out_features = out_features
        Conv = DWConv if depthwise else BaseConv

        # -----------------------------------------------#
        #   输入图片是 640（w）, 640（h）, 3（c）
        #   初始的基本通道是 64
        # -----------------------------------------------#
        base_channels = int(wid_mul * 64)  # 64
        base_depth = max(round(dep_mul * 3), 1)  # 3

        # -----------------------------------------------#
        #   改进：引入 SimAM 注意力机制
        # -----------------------------------------------#
        self.simam1 = SimAM(base_channels * 4)    # 256
        self.simam2 = SimAM(base_channels * 8)    # 512
        self.simam3 = SimAM(base_channels * 16)   # 1024

        # stem
        # -----------------------------------------------#
        #   利用 focus 网络结构进行特征提取
        #   640, 640, 3 -> 320, 320, 12 -> 320, 320, 64
        # -----------------------------------------------#
        self.stem = Focus(3, base_channels, ksize=3, act=act)

        # dark2
        # ---------------------------------------------------#
        #   完成卷积之后，320, 320, 64 -> 160, 160, 128
        #   完成 CSPlayer 之后，160, 160, 128 -> 160, 160, 128
        # ---------------------------------------------------#
        self.dark2 = nn.Sequential(
            Conv(base_channels, base_channels * 2, 3, 2, act=act),
            CSPLayer(
                base_channels * 2,  # 128
                base_channels * 2,  # 128
                n=base_depth,
                depthwise=depthwise,
                act=act,
            ),
        )

        # dark3
        # -----------------------------------------------#
        #   完成卷积之后，160, 160, 128 -> 80, 80, 256
        #   完成 CSPlayer 之后，80, 80, 256 -> 80, 80, 256
        # -----------------------------------------------#
        self.dark3 = nn.Sequential(
            Conv(base_channels * 2, base_channels * 4, 3, 2, act=act),
            CSPLayer(
                base_channels * 4,  # 256
                base_channels * 4,  # 256
                n=base_depth * 3,
                depthwise=depthwise,
                act=act,
            ),
        )

        # dark4
        # -----------------------------------------------#
        #   完成卷积之后，80, 80, 256 -> 40, 40, 512
        #   完成 CSPlayer 之后，40, 40, 512 -> 40, 40, 512
        # -----------------------------------------------#
        self.dark4 = nn.Sequential(
            Conv(base_channels * 4, base_channels * 8, 3, 2, act=act),
            CSPLayer(
                base_channels * 8,  # 512
                base_channels * 8,  # 512
                n=base_depth * 3,
                depthwise=depthwise,
                act=act,
            ),
        )

        # dark5
        # -------------------------------------------------#
        #   完成卷积之后，40, 40, 512 -> 20, 20, 1024
        #   完成 SPP 之后，20, 20, 1024 -> 20, 20, 1024
        #   完成 CSPlayer 之后，20, 20, 1024 -> 20, 20, 1024
        # -------------------------------------------------#
        self.dark5 = nn.Sequential(
            Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),
            SPPBottleneck(base_channels * 16, base_channels * 16, activation=act),
            CSPLayer(
                base_channels * 16,  # 1024
                base_channels * 16,  # 1024
                n=base_depth,
                shortcut=False,
                depthwise=depthwise,
                act=act,
            ),
        )

    def forward(self, x):
        outputs = {}
        x = self.stem(x)
        outputs["stem"] = x
        x = self.dark2(x)
        outputs["dark2"] = x
        # ------------------------------------------------------------#
        #   dark3的输出为 80, 80, 256，是一个有效特征层，引入 SimAM 模块
        # ------------------------------------------------------------#
        x = self.dark3(x)
        x1 = self.simam1(x)
        outputs["dark3"] = x1
        # ------------------------------------------------------------#
        #   dark3的输出为 40, 40, 512，是一个有效特征层，引入 SimAM 模块
        # ------------------------------------------------------------#
        x = self.dark4(x)
        x2 = self.simam2(x)
        outputs["dark4"] = x2
        # -------------------------------------------------------------#
        #   dark3的输出为 20, 20, 1024，是一个有效特征层，引入 SimAM 模块
        # -------------------------------------------------------------#
        x = self.dark5(x)
        x3 = self.simam3(x)
        outputs["dark5"] = x3
        return {k: v for k, v in outputs.items() if k in self.out_features}

Backbone 添加 SimAM 注意力机制之后的网络结构图如下。

  

5. 训练代码

定位到 YOLOX-main/tool/ 目录下的 train.py 文件，修改相应配置参数后右键运行此文件即可。

  欢迎大家订阅我的专栏一起学习 YOLO！（o^^o）