深入探索EPSA：提升卷积神经网络性能的新式注意力模块

 原论文地址：https://arxiv.org/abs/2105.14447

摘要

摘要部分提出了一种新的注意力模块——金字塔分割注意力（PSA）模块，该模块通过替代ResNet瓶颈块中的3x3卷积，显著提升了模型性能。PSA模块能够作为即插即用组件，增强网络的多尺度表征能力，使EPSANet在多个计算机视觉任务上超越了现有的通道注意力方法。

引言

引言部分讨论了注意力机制在计算机视觉领域的广泛应用，包括图像分类、目标检测等，并指出现有方法在处理空间信息和建立长距离通道依赖性方面的不足。作者强调了开发一种新的低成本且有效注意力模块的必要性，以改善这些问题。

总结

结论部分总结了PSA模块的优势，它能够提取丰富的多尺度空间信息，并在通道注意力向量中捕捉跨维度的重要特征。EPSA块通过增强多尺度表征能力和长距离通道依赖性，提高了网络性能。EPSANet在各种计算机视觉任务上都取得了优异的成绩，证明了其有效性。作者还提出了将PSA模块集成到更多轻量级CNN架构中的未来研究方向。

全局要点

1. 创新的PSA模块：

   • 多尺度特征提取：PSA模块通过将输入特征图分割成多个具有不同通道数的组，并在每个组上应用不同尺度的卷积核，实现了对多尺度空间信息的有效捕捉。这种分割和并行处理策略，使得网络能够在不同尺度上独立学习特征，增强了对局部和全局上下文的理解。
   • 长距离通道依赖性：通过Softmax操作对通道注意力权重进行重新校准，PSA模块能够在不同通道间建立长距离的依赖关系，这有助于网络更好地理解和整合全局信息，从而提升特征的表达能力。

2. EPSANet架构的性能提升：

   • 即插即用组件：EPSANet通过将PSA模块集成到ResNet的瓶颈块中，形成了一种新的网络架构。这种设计使得EPSA块可以作为一个即插即用的组件，轻松地融入到现有的网络中，无需对网络的其他部分进行大规模的修改。
   • 显著的性能增益：在图像分类、目标检测和实例分割等多个计算机视觉任务上，EPSANet展现出了卓越的性能。特别是在ImageNet数据集上的实验结果表明，EPSANet在Top-1准确率上超越了SENet等现有最先进的方法，证明了其强大的表征能力和优越的性能。
3. 网络结构的优化与效率：
   • 灵活的网络设计：EPSANet的网络结构设计允许通过调整PSA模块中的内核和组大小，来平衡模型的复杂度和性能。这种灵活性使得EPSANet能够适应不同的计算资源和应用需求。
   • 多尺度特征融合：EPSANet通过其独特的SPC（Split and Concat）模块，实现了在不同尺度上的特征融合，这有助于网络更全面地捕捉图像的细节和上下文信息，从而提高分类和检测的准确性。
   • 参数和计算效率：尽管EPSANet在性能上取得了显著提升，但其设计注重参数数量和计算成本的优化。与SENet等其他先进模型相比，EPSANet在保持高精度的同时，减少了参数量和计算量，提高了效率。
   • 扩展性和适应性：EPSANet的模块化设计不仅使得网络易于扩展，也便于根据不同的视觉任务调整网络结构。这种设计哲学为未来在更广泛的应用场景中部署和优化EPSANet提供了便利。

EPSANet，一种集成了金字塔分割注意力（PSA）模块的新型卷积神经网络架构，该模块通过多尺度特征提取和跨维度的长距离通道依赖性，显著提升了网络在图像分类、目标检测和实例分割等计算机视觉任务的性能。EPSANet的设计在保持高效计算的同时，实现了对特征的深入挖掘和有效利用，通过广泛的实验验证了其有效性，并开源了代码以促进社区的进一步研究和开发，同时展望了将PSA模块应用于更广泛网络架构的未来研究方向。

 PSA模块介绍

PSA模块是一种创新的注意力机制，通过在通道维度上分割特征图并应用多尺度金字塔卷积结构，捕获不同尺度的空间信息，然后利用Softmax操作重新校准通道注意力权重，以增强特征表示并建立长距离的通道依赖性，最终通过元素级乘法和特征融合提升模型的多尺度表征能力。

1. SEWeight模块：这是Squeeze-and-Excitation（SE）模块的一个变体，用于提取通道注意力权重。它通过全局平均池化（Global Average Pooling, GAP）获取输入特征图的通道统计信息，然后通过全连接层和激活函数（如ReLU和Sigmoid）来生成每个通道的注意力权重。

2. Softmax：Softmax函数用于将SEWeight模块输出的原始注意力权重转换为可以进行加权求和的归一化权重。Softmax确保了所有通道的注意力权重加起来等于1，从而可以对不同通道的贡献进行加权。

3. 元素级乘法（element-wise product）：在Softmax归一化之后，每个通道的注意力权重会与相应的特征图进行元素级乘法操作。这种操作允许网络在保持特征图的空间结构的同时，对不同通道的特征响应进行加权，强化重要的特征并抑制不重要的特征。

4. SPC模块（Split and Concatenate）：这是PSA模块的核心部分，负责实现多尺度特征提取。输入特征图首先在通道维度上被分割成S组，每组具有C/S个通道。然后，每组特征图通过不同尺度的卷积核进行处理，以捕获不同尺度的空间信息。处理后的多尺度特征图再次在通道维度上进行拼接，形成一个综合了多尺度信息的特征图。

5. 输入和输出：PSA模块的输入是一个高维的特征图（H×W×C），其中H和W分别是特征图的高度和宽度，C是通道数。经过PSA模块处理后，输出的也是一个高维特征图，其维度与输入相同，但是包含了经过注意力加权的多尺度空间信息。

整体来看，PSA模块的设计目的是通过在通道维度上的分割、多尺度卷积处理和注意力权重的重新校准，来增强特征图的表征能力，从而提高整个网络的性能。

SPC模块

这张图描述的是PSA（Pyramid Split Attention）模块中的一个关键组成部分，称为SPC（Split and Concatenate）模块。以下是对图中内容的分析：

1. Split（分割）：输入特征图X0X0​首先在通道维度上被分割成多个子特征图，这里表示为X0,X1,X2,X3X0​,X1​,X2​,X3​等。每个XiXi​是原始特征图的一个子集，具有C′=C/SC′=C/S个通道，其中CC是原始特征图的通道数，SS是分割成的组数。

2. Conv（卷积）：每个分割后的特征图XiXi​通过不同尺寸的卷积核进行卷积操作。图中显示了三种不同尺寸的卷积核，分别为K0×K0K0​×K0​，K1×K1K1​×K1​，和K2×K2K2​×K2​，以及对应的组大小G0,G1,G2G0​,G1​,G2​。这些卷积操作用于提取不同尺度的特征，从而捕获输入数据的多尺度空间信息。

3. KoxKo,GoKoxKo,Go：这可能表示第一个卷积层的内核大小和组大小，但具体数值未给出。通常，这些参数决定了卷积操作的感受野和计算复杂度。

4. Fo, F1, F2, F3：这些表示经过卷积操作后得到的特征图。每个FiFi​对应于一个分割的特征图XiXi​经过卷积处理的结果。

5. Concat（拼接）：所有经过卷积操作的特征图FiFi​在通道维度上进行拼接，形成一个综合了多尺度信息的单一特征图。这个拼接操作允许网络在后续处理中同时考虑不同尺度的特征。

6. 输出：最终，拼接后的特征图将作为SPC模块的输出，它包含了来自不同卷积层的多尺度特征信息，这些信息对于提高模型的表征能力是有益的。

整体来看，SPC模块是PSA中用于实现多尺度特征提取的关键步骤，它通过分割、卷积和拼接操作，有效地整合了来自不同尺度的特征，为后续的注意力加权和特征融合提供了丰富的信息。这种设计有助于提升模型对不同尺度特征的捕捉能力，从而增强其在复杂视觉任务中的性能。

代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class SPCModule(nn.Module):
    def __init__(self, channels, kernel_sizes, groups):
        super(SPCModule, self).__init__()
        self.splits = nn.ModuleList()
        self.kernel_sizes = kernel_sizes
        self.groups = groups
        for i, (kernel_size, group) in enumerate(zip(kernel_sizes, groups)):
            self.splits.append(self._make_split_conv(channels, kernel_size, group))

    def _make_split_conv(self, in_channels, kernel_size, groups):
        mid_channels = in_channels // len(self.kernel_sizes)
        layers = [nn.Conv2d(in_channels, mid_channels, kernel_size, groups=groups, padding=kernel_size//2)]
        layers.extend([
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, mid_channels, 1, groups=groups)
        ])
        layers.append(nn.ReLU(inplace=True))
        return nn.Sequential(*layers)

    def forward(self, x):
        split_features = [split_conv(x) for split_conv in self.splits]
        return torch.cat(split_features, dim=1)

class SEWeightModule(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SEWeightModule, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        avg_out = self.fc(self.avg_pool(x).view(b, -1))
        out = x * avg_out.expand_as(x)
        return out

class PSAModule(nn.Module):
    def __init__(self, in_channels, kernel_sizes=[3, 5, 7], groups=[1, 4, 8], reduction=16):
        super(PSAModule, self).__init__()
        self.spc = SPCModule(in_channels, kernel_sizes, groups)
        self.se = SEWeightModule(in_channels, reduction)

    def forward(self, x):
        x = self.spc(x)
        x = self.se(x)
        return x


psa_module = PSAModule(in_channels=64, kernel_sizes=[3, 5, 7], groups=[1, 4, 8])
# Assuming 'input_tensor' is the input feature map with shape (batch_size, 64, height, width)
output = psa_module(input_tensor)

这段代码是一个PyTorch实现的PSA（Pyramid Split Attention）模块，它由几个关键部分组成：

1. SPCModule 类：这个类实现了Split and Concatenate（SPC）模块，它负责多尺度特征的提取。

   • __init__ 方法中定义了多个卷积层，每个卷积层对应一个不同的内核尺寸（kernel_sizes），并且每个卷积层的组（groups）也不同。
   • _make_split_conv 是一个辅助方法，用于创建具有ReLU激活函数和1x1卷积的卷积层，这些层用于在不同的组上应用不同尺寸的卷积核。
   • forward 方法将输入特征图通过所有定义的卷积层，并将结果在通道维度上进行拼接。

2. SEWeightModule 类：这个类实现了SE（Squeeze-and-Excitation）权重模块，它用于计算通道注意力权重。

   • __init__ 方法中定义了全局平均池化和两层全连接层，其中包含了ReLU激活函数和Sigmoid激活函数，用于生成每个通道的权重。
   • forward 方法首先通过全局平均池化和全连接层计算出权重，然后将这些权重应用于输入特征图上，实现通道的加权。

3. PSAModule 类：这个类将SPCModule和SEWeightModule结合起来，实现了完整的PSA模块。

   • __init__ 方法中初始化了SPC模块和SE权重模块，允许用户指定输入通道数、内核尺寸和组大小。
   • forward 方法首先将输入特征图通过SPC模块提取多尺度特征，然后将结果通过SE权重模块进行加权。

4. 示例用法：在代码的最后部分，展示了如何实例化PSA模块，并将其应用于输入特征图。

整体来看，这段代码实现了PSA模块的核心功能，即通过多尺度卷积和通道注意力机制来增强特征图的表征能力。代码结构清晰，模块化设计使得PSA模块可以容易地集成到更大的网络架构中。不过，实际使用时可能需要根据具体的网络配置和参数进行调整。