SA-NET: Shuffle attention for DCNN 论文学习

Abstract

注意力机制可以使网络准确地关注到与输入相关的信息上，已成为深度神经网络提升性能的重要组成。在计算机视觉领域，主要有两种注意力机制，空间注意力机制和通道注意力机制，分别去获取像素对之间的关系和通道依赖关系。尽管将二者结合起来使用可能取得更优的表现，但是会不可避免地增加算力消耗。本文中，作者提出了一个高效率的 Shuffle Attention 模块来解决此问题，它采用 Shuffle Units 来将两种注意力机制有效地结合起来。SA 首先将通道维度分组为多个亚特征，然后并行地处理它们。然后对于每个亚特征，SA 利用一个 Shuffle Unit 来描述其在空间和通道维度上的特征依赖关系。然后，所有的亚特征会被聚集起来，使用一个“channel shuffle”来让信息在不同的亚特征之间传递。该SA模块的效率很高，比如，SA 和主干网络 ResNet50 的参数量和计算量分别是：300 vs. 25.56M，$2.76e-3$ GFLOPS vs. $4.12$ GFLOPS，而 top-1 准确率提升了 $1.34\%$。作者在分类数据集 ImageNet-1K 和目标检测数据集 MS COCO 以及实例分割任务上做了实验，表明 SA 的性能要超过目前 SOTA 的方法，实现了更高的准确率，而模型复杂度较低。代码和模型都放在了 https://github.com/wofmanaf/SA-Net。

1. Introduction

注意力机制在研究领域吸引了越来越多的关注，因为它可以提升特征表示，关注在重要特征上，而抑制不重要的那些。最近的研究表明，在卷积模块中正确地融入注意力机制能够极大地提高视觉任务的表现，如图像分类、目标检测和实例分割。

在计算机视觉领域中，主要有两类注意力机制：通道注意力机制和空间注意力机制，通过不同的聚合策略、变换和增强函数，将所有位置的相同特征聚合起来。基于这些观察，GCNet 和 CBAM 在一个模块中集成了空间注意力和通道注意力，实现了极大的提升。但是，它们都很难收敛，且计算量很大。另一些研究则试图简化通道或空间注意力的结构。例如，ECA-Net 通过一维卷积简化了 SE 模块的通道权重计算。SGE 将通道维度分组为亚特征，来代表不同的语义信息，通过一个注意力mask对所有位置的特征向量做缩放，给每个特征组都使用了一个空间机制。但是，它们没有充分利用空间和通道注意力的关系，所以就没那么高效率。“我们能否将不同的注意力模块融合得更轻量而高效呢？”

图1。最新 SOTA 注意力模块（SENet、CBAM、ECA-Net、SGE-Net、SA-Net）在 ImageNet-1K 上的比较，指标有准确率、参数量和计算量GFLOPs。圆圈大小代表 GFLOPs。很明显，SA-Net 的准确率更高、模型复杂度更低。

为了回答这个问题，作者首先回顾了下 ShuffleNet v2 中的 unit，它能高效率地构建一个多分支结构，并行地处理不同的分支。在每个 unit 的开始阶段，$c$个特征通道的输入被划分为2个分支，通道数分别是$c-c^{\prime }$和$c^{\prime }$。随后，采用多个卷积层来获取更高层级的输入特征表示。然后将2个分支 concat 起来，使通道数与输入的通道数相同。最后，采用 channel shuffle 操作让两个分支的信息可以交流。此外，为了增加计算速度，SGE 引入了一个分组策略，将输入特征图沿着特征维度划分为若干个组。然后就可并行地增强所有的亚特征。

基于这些发现，本文提出了一个轻量且高效率的 SA 模块，将通道维度分组为亚特征。对于每个亚特征，SA 采用 Shuffle Unit 来同时构建通道注意力和空间注意力。对于每个注意力模块，本文设计了一个覆盖所有位置的注意力 mask，抑制可能的噪音，凸显正确的语义特征区域。实验结果显示（如图1）表明，该模块简单且高效，与现有方法相比，参数量少，准确率高。

本文核心贡献有：1) 引入了一个轻量但高效的注意力模块，SA，将通道维度分组为多个亚特征，然后对每个亚特征，利用 Shuffle Unit 来集成通道和空间注意力模块。2) 在 ImageNet-1k 和 MS COCO 上的实验表明，SA 的模型复杂度要低于现有 SOTA 方法，性能却高于它们。

2. Related Work

多分支结构。CNN 多分支结构在准确性和速度方面都取得了进步。多分支结构的原理就是“拆分-变换-融合”，降低网络的训练难度。InceptionNet 系列就是个成功的多分支结构，每个分支都精心地设计了卷积核滤波器，来融合更丰富的特征。ResNets 也可看作为一个双分支网络，一个分支是恒等映射。SKNet 和 ShuffleNet 家族都延续了 InceptionNet 的思想，但它们至少在两个方面是不同的。SKNets 利用了一个自适应的选择机制来实现神经元自适应的感受野大小。ShuffleNets 进一步将“通道拆分”和“通道 shuffle”操作融合进了一个操作里面，在速度和精度上取得平衡。

分组特征。将学习特征分为多个组，可以追溯到 AlexNet，目的是将模型分配到更多的GPU资源上。Deep Roots 测试了 AlexNet，指出卷积分组能够学到更优的特征表示。MobileNets 和 ShuffleNets 将每个通道看作为一个组，在这些组内进行空间关系建模。CapsuleNets 将每个神经元组看作为一个胶囊，活跃胶囊里的神经元活动表示图像中某实例的各种属性。SGE 设计了 CapsuleNets，将通道维度划分为多个亚特征来学习不同的语义。

注意力机制。在很多论文中都研究了注意力的重要性。它强调信息丰富的特征表示，抑制用处不大的特征。自注意力方法计算一个位置上的上下文信息，作为图像所有位置的加权和。SE 使用2个全连接层来对元素间关系建模。ECA-Net 采用一个一维卷积滤波器来生成通道权重，极大地降低 SE 的模型复杂度。Wang 等人提出了 non-local 模块，计算特征图上每个空间点之间的关系矩阵，生成一个注意力图。CBAM、GCNet和SGE 将空间注意力和通道注意力结合起来，而 DANet 对不同分支的两个注意力模块求和，自适应地将局部特征和全局依赖关系集成起来。

3. Shuffle Attention

本章节，作者首先介绍 SA 模块构建的过程，将输入特征图拆分为多个组，然后对每个组使用 Shuffle Unit，将通道注意力和空间注意力集成入一个模块中去。然后，聚合所有的亚特征，使用"channel shuffle"操作让信息在不同的亚特征之间传递。然后介绍如何在 CNN 中使用 SA。最后，作者可视化了SA的效果，验证其可靠性。图2就是 SA 模块的整体架构。

特征分组。对于给定的特征图 $X\in {\mathbb{R}}^{C\times H\times W}$，其中$C,H,W$分别是通道数、空间高度和宽度。SA 首先沿着通道维度将 $X$ 划分为 $G$个组，即 $X=[X_1,...,X_G],X_k\in {\mathbb{R}}^{C/G\times H\times W}$，在训练过程中每个亚特征$X_k$逐渐地获取一个语义响应。然后通过注意力模块，我们为每个亚特征都生成一个相应的重要度系数。在每个注意力 unit 的开始时刻，$X_k$的输入会沿着通道维度，被分为2个分支，即$X_{k1},X_{k2}\in {\mathbb{R}}^{C/2G\times H\times W}$。如图2，一个分支利用通道间的相互关系，输出通道注意力图，另一个分支则利用特征的空间关系，生成空间注意力图。这样，模型就可以关注在那些有意义的信息上。

图2. SA 模块的整体结构。它采用“channel split”并行地处理每个组的亚特征。对于通道注意力分支，通过 GAP 来产生通道统计数据，然后用一对参数来缩放并平移通道向量。对于空间注意力分支，采用 group norm 来产生空间统计数据，然后创建出一个紧致的向量，与通道分支类似。然后将这两个分支 concat 起来。随后聚合所有的亚特征，最终利用“channel shuffle”操作来使信息在不同的亚特征之间传递。

通道注意力。充分获取通道依赖关系的一个方式就是 SE 模块。但是，它所使用的参数过多，要想设计一个轻量级注意力模块并不太好。同样，类似于 ECA 中的通过一个更快速的、大小是$k$的一维卷积操作来生成通道权重也不适合，因为$k$通常会比较大。因此作者提供了一个替代项，它首先通过全局平均池化来生成通道统计数据，$s\in {\mathbb{R}}^{C/2G\times 1\times 1}$，嵌入全局信息，可以沿着空间维度$H\times W$收缩$X_{k1}$计算得到。

$$s={\mathcal{F}}_{gp}(X_{k1})=\frac{1}{H\times W}\sum _{i=1}^H\sum _{j=1}^W{X}_{k1}(i,j)$$

此外，会创建一个紧致的特征来准确地、自适应地选择。这是通过一个简单的门机制实现的，用 sigmoid 激活函数做到。通道注意力的最终输出就是：

$$X_{k1}^{\prime }=\sigma ({\mathcal{F}}_c(s))\cdot X_{k1}=\sigma (W_{1}s+b_1)\cdot X_{k1}$$

其中$W_1\in {\mathbb{R}}^{C/2G\times 1\times 1}$和$b_1\in {\mathbb{R}}^{C/2G\times 1\times 1}$用于缩放和平移 $s$。

空间注意力。与通道注意力不同，空间注意力关注在“哪里”是信息丰富的，是通道注意力的补充。首先，作者对$x_{k2}$使用 Group Norm，获取空间统计数据。然后使用$F_c(\cdot )$来增强${\hat{X}}_{k2}$的特征表示。最终的空间注意力输出是：

$$X_{k2}^{\prime }=\sigma (W_2\cdot GN(X_{k2})+b_2)\cdot X_{k2}$$

其中$W_2,b_2$ 是形状为${\mathbb{R}}^{C/2G\times 1\times 1}$的参数。

然后两个分支的结果会被 concat 起来，使得通道个数与输入的个数相同，即$X_k^{\prime }=[X_{k1}^{\prime },X_{k2}^{\prime }]\in {\mathbb{R}}^{C/G\times H\times W}$。

聚合。随后，所有的特征会被聚合起来。最终与 ShuffleNet v2 相似，作者采用了 channel shuffle 的操作，沿着通道维度实现跨组信息交流。SA 模块的最终输出与$X$的形状相同，使得 SA 模块可以很容易地集成到其它网络中去。

注意$W_1,b_1,W_2,b_2$ 和 Group Norm 的超参数只是 SA 模块中的参数。在单个 SA 模块中，每个分支通道的个数是 $C/2G$。因此，所有的参数个数就是$3C/G$，通常$G$是32或64，这和网络数以百万计的参数相比微乎其微，SA 就非常地轻量。SA 整体架构在图2中有展示。

实现。SA 可以通过 PyTorch 和 TensorFlow 很容易就实现。图3就是 PyTorch 的代码。

SA-Net。在 CNN 中加入 SA 模块，作者使用了与 SENet 相同的配置，将 SA 模块替换为 SE 模块。该网络的名称就是 SA-Net。

3.1 可视化和理解

为了验证 SA 能否通过特征分组和 channel shuffle 来提升语义特征表示，作者首先在 ImageNet-1k 上训练了 SA-Net50B（没用 channel shuffle）和 SA-Net50（用了 channel shuffle）。假设 $I$ 是原始输入，作者计算了SA模块用之前和用之后，每个组内SA_5_3位置（SA_stageID_blockID中最后一个 bottleneck）上 $I\times X_k$ 的 top-1 准确率。作者将准确率得分作为指标，在图4（a,b）中画出了ImageNet-1k验证集内不同类别（狮子狗、金鱼、飞机）的分布。为了比较，作者也画出了所有1000个类别的分类得分的分布情况。

图4. SA 有效性的验证

如图4(a,b)，在应用了 SA 之后，top-1 准确率有提高，也就是说特征分组可以明显增强特征图的语义表示。此外在使用了 channel shuffle之后，每个分组的平均得分都增长了 $\approx 0.4\%$，证明了 channel shuffle 的有效性。

为了充分验证 SA 的有效性，作者将SA-Net50（用了 shuffle）在不同深度的三个类别（狮子狗、金鱼、飞机）的平均激活（与 SE 相似，每个组内通道特征图的均值）分布画了出来。结果在图4©中展示。关于SA模块，作者做了一些观察：(1) 在早期网络层，不同类别的分布非常相似，（比如 SA_2_3和SA_3_4），这就说明在网络早期阶段，不同类别可能会共享特征组的重要度；(2) 网络越深，每个组的激活就变得越来越和类别相关，不同的类别具有不同的特征判别值，表现也就不同（如SA_4_6和SA_5_3）。(3) SA_5_2在不同的类别上表现了相似的规律，这就说明与其它模块相比，SA_5_2对于给网络提供重新校准没那么重要。

为了验证 SA 的作用，作者从 ImageNet-1k 验证集中选择了9张图片。使用了 Grad-CAM 来可视化其在 SA-Net50上 SA_5_3 的热力图。为了比较，作者也画了 ResNet50 在"layer4.2"的热力图。如图5所示，SA模块使分类模型更加关注在相关的区域，有更多的目标细节信息，也就是说 SA 模块能够有效地提升分类准确率。

因此，SA模块能够实际增强网络的特征表示能力。