即插即用模块-EMA跨空间学习的高效多尺度注意模块（论文+代码）

松小白song

已于 2024-05-23 15:38:12 修改

阅读量2.1w

点赞数 54

分类专栏：即插即用模块文章标签：学习

于 2024-05-23 15:09:07 首次发布

本文链接：https://blog.csdn.net/weixin_47151388/article/details/139147329

版权

即插即用模块专栏收录该内容

5 篇文章

订阅专栏

一、摘要

二、模型概要

2.1Revisit Coordinate Attention (CA)

2.2Multi-Scale Attention (EMA) Module

论文：https://arxiv.org/pdf/2305.13563v2

代码：https://github.com/yoloonme/ema-attention-module

一、摘要

在各种计算机视觉任务中，通道或空间注意机制对于产生更多可识别的特征表示具有显着的有效性。然而，通过通道降维来建模跨通道关系可能会对提取深度视觉表征带来副作用。本文提出了一种新型的高效多尺度注意力(EMA)模块。为了保留每个通道上的信息和减少计算开销，我们将部分通道重构为批处理维度，并将通道维度分组为多个子特征，使空间语义特征在每个特征组内均匀分布。具体而言，除了编码全局信息以重新校准每个并行分支中的通道权重外，还通过跨维交互进一步聚合两个并行分支的输出特征，以捕获像素级成对关系。我们使用流行的基准(例如CIFAR-100、ImageNet-1k、MS COCO和VisDrone2019)对图像分类和目标检测任务进行了广泛的烧烧研究和实验，以评估其性能。

对比以ResNet50为骨架的不同关注方法的准确率，可以发现EMA表现出最高的Top-1准确率，且模型复杂度较低。模型也在COCO基准上以yolov5s (v6.0)为标准主干进行了评估，这说明了EMA是高效且有效的。

我们提出了一种新的跨空间学习方法，并设计了一个多尺度并行子网络来建立短期和长期依赖关系。

主要贡献如下:
1.提出了一种新的跨空间学习方法，并设计了一个用于建立短期和长期依赖关系的多尺度并行子网络；

2.考虑了一种通用方法，将部分通道维度重塑为Batch维度，以避免通过通用卷积进行某种形式的降维；

3.除了在没有通道降维的情况下在每个并行子网络中建立局部跨通道交互外，作者还通过跨空间学习方法融合了两个并行子网络的输出特征图；
4.与CBAM、基于归一化的注意力模块(NAM)、SA、ECA和CA相比，EMA不仅取得了更好的结果，而且在所需参数方面更高效。

二、模型概要

其中，“g”表示分组，“X Avg Pool”表示一维水平全局池，“Y Avg Pool”表示一维垂直全局池。

2.1Revisit Coordinate Attention (CA)

CA模块与SE注意力模块相似，通过全局平均池化操作来建模跨通道信息。具体来说，CA模块将输入张量分解为两个并行的一维特征编码向量，以建模带有空间位置信息的跨通道依赖关系。通过在水平维度方向上的一维全局平均池化，CA捕获了沿垂直维度方向的全局信息。类似地，在垂直维度方向上的一维全局平均池化捕获了沿水平维度方向的全局信息。这些处理使CA不仅能够保持精确的位置信息，还能够有效地利用长程依赖。

2.2Multi-Scale Attention (EMA) Module

EMA模块通过并行子结构避免了更多的顺序处理和更大的深度。EMA的整体结构如图3(b)所示。以下讨论了EMA如何在卷积操作中不进行通道维度减缩的情况下学习有效的通道描述，并为高级特征图生成更好的像素级注意力。

具体实现方面，EMA采用了CA模块中的共享1x1卷积分支，并将其命名为1x1分支。此外，为了聚合多尺度的空间结构信息，EMA在1x1分支并行放置了一个3x3卷积核，命名为3x3分支。考虑到特征分组和多尺度结构，EMA能够有效地建立短程和长程依赖，从而提高性能。

特征分组

对于任何给定的输入特征图，EMA将 𝑋 分为 𝐺 个子特征组，每个组学习不同的语义。特征分组方法使得模型能够在更多的GPU资源上分配和处理模型。这种分组方式不仅加强了语义区域的特征学习，还压缩了噪声。

并行子网络

EMA采用三个并行路径来提取分组特征图的注意力权重描述符。两个路径为1x1分支，第三个路径为3x3分支。1x1分支中采用了一维全局平均池化操作，以分别在两个空间方向上编码通道信息。3x3分支则通过3x3卷积捕获多尺度特征表示。这样，EMA不仅能够编码跨通道的信息，以调整不同通道的重要性，还能够将精确的空间结构信息保留到通道中。

跨空间学习

EMA通过在不同空间维度方向上提供跨空间信息聚合方法来丰富特征聚合。具体来说，1x1分支的输出通过二维全局平均池化编码全局空间信息，而3x3分支的输出直接转换为相应的维度形状。然后，通过矩阵点积操作聚合这些输出，生成第一个空间注意力图。最终，每个组内的输出特征图通过两个生成的空间注意力权重值的Sigmoid函数进行聚合，捕获像素级配对关系，并突出所有像素的全局上下文。EMA的最终输出与输入大小相同，可以有效地堆叠到现代架构中。

2.3表格对比CA和EMA

三、代码解读

import torch  # 导入 PyTorch 库
from torch import nn  # 从 PyTorch 中导入神经网络模块

class EMA(nn.Module):  # 定义一个继承自 nn.Module 的 EMA 类
    def __init__(self, channels, c2=None, factor=32):  # 构造函数，初始化对象
        super(EMA, self).__init__()  # 调用父类的构造函数
        self.groups = factor  # 定义组的数量为 factor，默认值为 32
        assert channels // self.groups > 0  # 确保通道数可以被组数整除
        self.softmax = nn.Softmax(-1)  # 定义 softmax 层，用于最后一个维度
        self.agp = nn.AdaptiveAvgPool2d((1, 1))  # 定义自适应平均池化层，输出大小为 1x1
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))  # 定义自适应平均池化层，只在宽度上池化
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))  # 定义自适应平均池化层，只在高度上池化
        self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)  # 定义组归一化层
        self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)  # 定义 1x1 卷积层
        self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)  # 定义 3x3 卷积层

    def forward(self, x):  # 定义前向传播函数
        b, c, h, w = x.size()  # 获取输入张量的大小：批次、通道、高度和宽度
        group_x = x.reshape(b * self.groups, -1, h, w)  # 将输入张量重新形状为 (b * 组数, c // 组数, 高度, 宽度)
        x_h = self.pool_h(group_x)  # 在高度上进行池化
        x_w = self.pool_w(group_x).permute(0, 1, 3, 2)  # 在宽度上进行池化并交换维度
        hw = self.conv1x1(torch.cat([x_h, x_w], dim=2))  # 将池化结果拼接并通过 1x1 卷积层
        x_h, x_w = torch.split(hw, [h, w], dim=2)  # 将卷积结果按高度和宽度分割
        x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())  # 进行组归一化，并结合高度和宽度的激活结果
        x2 = self.conv3x3(group_x)  # 通过 3x3 卷积层
        x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))  # 对 x1 进行池化、形状变换、并应用 softmax
        x12 = x2.reshape(b * self.groups, c // self.groups, -1)  # 将 x2 重新形状为 (b * 组数, c // 组数, 高度 * 宽度)
        x21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))  # 对 x2 进行池化、形状变换、并应用 softmax
        x22 = x1.reshape(b * self.groups, c // self.groups, -1)  # 将 x1 重新形状为 (b * 组数, c // 组数, 高度 * 宽度)
        weights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)  # 计算权重
        return (group_x * weights.sigmoid()).reshape(b, c, h, w)  # 应用权重并将形状恢复为原始大小