YOLO11改进|注意力机制篇|引入矩形自校准模块RCM

一、【RCM】注意力机制

1.1【RCM】注意力介绍

下图是【RCM】的结构图，让我们简单分析一下运行过程和优势

RCA 模块（Residual Cross-Attention）:
目的: RCA 模块实现了一种空间注意力机制，能够捕获特征图中高度和宽度两个维度的信息交互。这可以增强网络对空间区域的关注，提升它对局部上下文的理解。
处理步骤:
空间卷积: 模块首先使用深度可分离卷积（dwconv_hw）在输入张量的高度和宽度上进行处理，这有助于捕捉局部的空间信息。
自适应池化: 模块分别对高度（pool_h）和宽度（pool_w）方向进行全局平均池化，以获得每个通道的高度和宽度的信息聚合结果。高度池化会生成形状为 [N, C, H, 1] 的张量，而宽度池化会生成形状为 [N, C, 1, W] 的张量。
Excite模块: 将高度和宽度的信息相加（x_gather = x_h + x_w），然后通过一系列卷积、归一化和激活操作（ReLU + Sigmoid）生成注意力权重。这个注意力机制学习到了高度和宽度的相关性，并将其用于调整输入特征。
残差连接: 最后，将经过注意力权重调整的特征与原始卷积结果进行逐元素乘法相结合，输出最终的增强特征。这种操作可以强化重要的空间区域。
优势:RCA 模块引入了全局高度和宽度的池化操作，使得模型能够捕捉长距离依赖关系。使用深度可分离卷积减少了计算成本，同时保持了良好的特征提取能力。注意力机制可以动态地调整特征图中不同空间位置的权重，从而提高模型对重要局部信息的捕捉能力。
RCM 模块（Residual Cross Mixer）:
目的: RCM 模块整合了多种模块化操作（包括RCA 和 MLP）用于特征的混合与增强，同时通过残差连接确保信息不会丢失。

处理步骤:

Token Mixer (RCA): RCM 模块首先调用 RCA 模块，使用其自带的注意力机制对输入特征图进行初步处理。这个步骤能够捕捉特征图的空间关系，并输出加权后的特征图。
归一化: 对经过RCA处理后的输出应用批归一化层，稳定模型的训练并改善收敛性。
MLP层: 调用 ConvMlp 模块，将输入的特征图通过两个1x1卷积进行通道间的线性变换，进一步提取特征。这相当于一个在卷积上下文中使用的多层感知机。
Layer Scale (Gamma): RCM 模块使用一个可学习的参数 gamma，其初始值非常小（1e-6），在更新过程中逐步优化。这种机制允许模型在残差连接上进行适度的权重调节，使得模型能够更加灵活地融合新旧信息。
残差连接: 将经过 DropPath 操作后的输出与初始输入进行相加，实现残差连接。这确保了信息流不会因过多的非线性变换而丢失，增强了模型的训练效果。
优势 :RCM 模块通过结合 RCA 和 MLP 操作，既能捕捉全局的空间关系，又能实现通道间的特征交互。Layer Scale 使得模型可以平衡新旧特征之间的融合，增加了模型的鲁棒性。DropPath 操作引入了随机深度的正则化，帮助模型在训练时进行更好的泛化。

1.2【RCM】核心代码

python 复制代码

from functools import partial

import torch
import torch.nn as nn

from timm.models.layers import DropPath
from timm.models.layers.helpers import to_2tuple


class ConvMlp(nn.Module):
    """ MLP using 1x1 convs that keeps spatial dims
    copied from timm: https://github.com/huggingface/pytorch-image-models/blob/v0.6.11/timm/models/layers/mlp.py
    """

    def __init__(
            self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU,
            norm_layer=None, bias=True, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = to_2tuple(bias)

        self.fc1 = nn.Conv2d(in_features, hidden_features, kernel_size=1, bias=bias[0])
        self.norm = norm_layer(hidden_features) if norm_layer else nn.Identity()
        self.act = act_layer()
        self.drop = nn.Dropout(drop)
        self.fc2 = nn.Conv2d(hidden_features, out_features, kernel_size=1, bias=bias[1])

    def forward(self, x):
        x = self.fc1(x)
        x = self.norm(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        return x


class RCA(nn.Module):
    def __init__(self, inp, kernel_size=1, ratio=1, band_kernel_size=11, dw_size=(1, 1), padding=(0, 0), stride=1,
                 square_kernel_size=2, relu=True):
        super(RCA, self).__init__()
        self.dwconv_hw = nn.Conv2d(inp, inp, square_kernel_size, padding=square_kernel_size // 2, groups=inp)
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        gc = inp // ratio
        self.excite = nn.Sequential(
            nn.Conv2d(inp, gc, kernel_size=(1, band_kernel_size), padding=(0, band_kernel_size // 2), groups=gc),
            nn.BatchNorm2d(gc),
            nn.ReLU(inplace=True),
            nn.Conv2d(gc, inp, kernel_size=(band_kernel_size, 1), padding=(band_kernel_size // 2, 0), groups=gc),
            nn.Sigmoid()
        )

    def sge(self, x):
        # [N, D, C, 1]
        x_h = self.pool_h(x)
        x_w = self.pool_w(x)
        x_gather = x_h + x_w  # .repeat(1,1,1,x_w.shape[-1])
        ge = self.excite(x_gather)  # [N, 1, C, 1]

        return ge

    def forward(self, x):
        loc = self.dwconv_hw(x)
        att = self.sge(x)
        out = att * loc
        return out


class RCM(nn.Module):
    """ MetaNeXtBlock Block
    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
        ls_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """

    def __init__(
            self,
            dim,
            token_mixer=RCA,
            norm_layer=nn.BatchNorm2d,
            mlp_layer=ConvMlp,
            mlp_ratio=2,
            act_layer=nn.GELU,
            ls_init_value=1e-6,
            drop_path=0.,
            dw_size=11,
            square_kernel_size=3,
            ratio=1,
    ):
        super().__init__()
        self.token_mixer = token_mixer(dim, band_kernel_size=dw_size, square_kernel_size=square_kernel_size,
                                       ratio=ratio)
        self.norm = norm_layer(dim)
        self.mlp = mlp_layer(dim, int(mlp_ratio * dim), act_layer=act_layer)
        self.gamma = nn.Parameter(ls_init_value * torch.ones(dim)) if ls_init_value else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        shortcut = x
        x = self.token_mixer(x)
        x = self.norm(x)
        x = self.mlp(x)
        if self.gamma is not None:
            x = x.mul(self.gamma.reshape(1, -1, 1, 1))
        x = self.drop_path(x) + shortcut
        return x


if __name__ == "__main__":
    # 创建随机输入张量 (batch_size, channels, height, width)
    x = torch.randn(8, 64, 32, 32)  # 假设输入是8个样本，64个通道，32x32大小

    # 实例化 ConvMlp 模块
    conv_mlp = ConvMlp(in_features=64, hidden_features=128, out_features=64, act_layer=nn.ReLU)

    # 实例化 RCA 模块
    rca = RCA(inp=64, kernel_size=1, ratio=1, band_kernel_size=11, square_kernel_size=3)

    # 实例化 RCM 模块
    rcm = RCM(dim=64, dw_size=11, square_kernel_size=3, ratio=1, drop_path=0.1)

    # 测试 ConvMlp 模块
    print("Testing ConvMlp module:")
    output_mlp = conv_mlp(x)
    print(f"Output shape of ConvMlp: {output_mlp.shape}")

    # 测试 RCA 模块
    print("\nTesting RCA module:")
    output_rca = rca(x)
    print(f"Output shape of RCA: {output_rca.shape}")

    # 测试 RCM 模块
    print("\nTesting RCM module:")
    output_rcm = rcm(x)
    print(f"Output shape of RCM: {output_rcm.shape}")

二、添加【RCM】注意力机制

2.1STEP1

首先找到ultralytics/nn文件路径下新建一个Add-module的python文件包【这里注意一定是python文件包，新建后会自动生成_init_.py】，如果已经跟着我的教程建立过一次了可以省略此步骤 ，随后新建一个RCM.py文件并将上文中提到的注意力机制的代码全部粘贴到此文件中，如下图所示

2.2STEP2

在STEP1中新建的_init_.py文件中导入增加改进模块的代码包如下图所示

2.3STEP3

找到ultralytics/nn文件夹中的task.py文件，在其中按照下图添加

2.4STEP4

定位到ultralytics/nn文件夹中的task.py文件中的def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)函数添加如图代码,【如果不好定位可以直接ctrl+f搜索定位】

三、yaml文件与运行

3.1yaml文件

以下是添加【RCM】注意力机制在小目标检测层后的yaml文件，大家可以注释自行调节，效果以自己的数据集结果为准