特征图注意力_CCNet：用于语义分割的交叉注意力网络

此篇文章内容源自 CCNet: Criss-Cross Attention for Semantic Segmentation，若侵犯版权，请告知本人删帖。

此篇文章是论文 CCNet: Criss-Cross Attention for Semantic Segmentation 的阅读笔记。

原论文下载地址：

CCNet: Criss-Cross Attention for Semantic Segmentation​arxiv.org

摘要

远程依赖可以捕获有用的上下文信息，从而有益于视觉理解问题。作者在此篇论文中提出了交叉网络(CCNet)，可以通过一种更有效的方法获取上述重要信息。具体地说，对每一个像素，CCNet可以通过一个新的交叉注意力模块获取该像素交叉路径上的邻近像素的上下文信息。通过进一步的递归操作，每个像素最终可以获取所有的像素的远程依赖。总体来看，CCNet有以下优点：1. 显存占用率合理。与非局部模块相比，其显存占用率是递归交叉注意力模块的11倍。2. 计算效率高。与非局部模块相比，在计算远程依赖时，递归交叉注意力明显地降低了85%的FLOPs。3. 最先进的性能。在主流的语义分割基准上进行了扩展实验，包括Cityscapes，ADE20K，以及实例分割基准COCO。特别地，CCNet在Cityscapes测试集和ADE20K验证集上达到了81.4和45.22的mIoU分值，是最新的最先进的结果。开源代码连接：speedinghzl/CCNet.

引言

为了生成密集的像素级上下文信息，PSANet^[1]通过一张预测的注意力图来学习聚合每个位置的上下文信息。非局部网络利用一个自注意力机制生成了更强大的像素级表示，其利用的自注意力网络可以使任意一个位置的单独特征能够感知其他所有的特征。至此，特征图中的每个位置通过自适应预测的注意力图连接到其他所有位置，因此获取了不同范围的上下文信息，如图1(a)所示。然而，这些基于注意力的方法需要生成巨大的注意力图来计算每个像素对之间的关系，其时间、空间复杂度均为

，

为输入特征图的空间维度。由于语义分割任务中的输入特征图通常具有高分辨率，所有基于自注意力的方法具有高计算复杂度，并且占用大量显存。

作者发现当前^[2]采用的非局部操作可以通过两个连续的交叉操作替代，每个交叉操作中，每个位置在特征图中仅有

个稀疏连接。这启发了作者提出交叉注意力模型，以此来聚合水平、垂直方向上的远程像素级上下文信息。通过串行堆叠两个交叉注意模块，它可以收集来自所有像素的上下文信息。该分解显著地将时间和空间复杂性从

降至

。

具体来说，交叉注意力模块可以获取交叉路径上，无论远近的各种信息。如图1所示，非局部模块和交叉注意力模块均使用空间尺寸为

的输入特征图来生成注意力图(上部分支)和与之适应的特征图(下部分支)。然后，采用加权求和作为聚合方法。在交叉注意力模块，特征图中的每个位置(如，蓝色标记)，通过预测的稀疏注意力图，都与其他位于同行同列的其他位置相连接。预测的注意力图只有

个权重，而不是非局部模块中的

个。此外，作者提出递归交叉注意力模块来获取所有像素的远程依赖。包含水平、垂直方向的上下文信息的局部特征仅进入交叉注意力模块一次。交叉注意力模块的输出特征图输入给下一个交叉注意力模块；第二个特征图中的每个位置(如，红色)，搜集其他位置的上下文信息来增强像素级的表示。所有交叉注意力模块共享参数，以此减少额外的参数。交叉注意力模块可以嵌入到任何全卷积神经网络，命名为CCNet，以此学习端到端的分割方式。

图 1 两种基于注意力的上下文聚合方法的关系图。

(a) 每个位置(如，蓝色标记)，非局部模块生成一个具有

个权重的密集注意力图(绿色标记)。(b) 每个位置(如，蓝色标记)，交叉注意力模块生成一个稀疏的注意力图，仅有

个权重。当迭代操作后，最终输出特征图中的每个位置(如，红色标记)，能够获取所有像素的远程依赖。为了便于展示，残余连接被略去了。

作者的主要贡献如下：

• 提出一个新的交叉注意力模块，可以一个更有效的方法从远程依赖中获取上下文信息。
• 利用两个迭代交叉注意力模型的优势，提出了CCNet，在分割任务基准上达到了领先的性能，分割任务基准包括Cityscapes，ADE20K，MSCOCO。

方法

本章节介绍CCNet的细节。首先，给出CCNet的总体框架。其次，介绍获取水平、垂直方向上远程上下文信息的交叉注意力模块。最后，为了获取密集的全局上下文信息，作者提出了递归交叉注意力模块。

方法-整体

网络框架如图2所示。一幅输入图像通过一个深度卷积神经网络(DCNN)送入网络，DCNN是按照全卷积方式设计的，生成一个特征图

。定义特征图

的尺寸为

。为了保留更多细节和有效地生成密集特征图，作者移除了最后两个下采样操作，并且在后续的卷积层中使用了扩张卷积，以此放大输出特征图

的宽和高至输入图像的

。

在获取了特征图

之后，首先使用一个卷积层获取降维的特征图

，然后将特征图

送入交叉注意力(CCA)模块，生成新的特征图

，特征图

以交叉的方式聚合了每个像素的远程上下文信息。特征图

仅聚合了水平、垂直方向的上下文信息，不足以用来进行语义分割。为了获取更丰富和密集的上下文信息，作者再次将特征图

送入交叉注意力模块，输出特征图

。因此，特征图

中每个位置实际上聚合了所有像素的信息。前后两个交叉注意力模块共享相同的参数，避免添加过多的额外参数。作者将这种递归结构命名为递归交叉注意力(RCCA)模块。

然后将密集的上下文特征图

与局部表示特征

相拼接。拼接之后，通过一个或多个具有BN和激活的卷积层进行特征融合。最后，将融合后的信息送入分割层来生成最后的分割图。

图 2 用于语义分割的CCNet的总体结构

递归交叉注意力将特征图

作为输入，输出特征图

，

包含了所有像素的丰富的、密集的上下文信息。递归交叉注意力模块可以展开为

的循环，在此循环中所有交叉注意力模块共享参数。

方法-交叉注意力

为了使用轻量级计算量和内存，在局部特征表示上建模远程上下文依赖，作者提出了交叉注意力模块。交叉注意力模块搜集水平、垂直方向上的上下文信息，以此加强像素级表示能力。

如图3所示，给定一个局部特征

，交叉注意力模块首先对

使用两个卷积核为

的卷积层，生成两个特征图

，其中

。

是特征图的通道数，为了降维，

小于

。

图 3 交叉注意力模块的细节

获得特征图

之后，通过

关联操作进一步生成注意力图

。在特征图

的空间维度中的每个位置

，可以得到一个向量

【各通道上位置

上的值构成的向量。】。同时，可以通过从特征图

抽取特征向量来获取集合

，集合

中的

对应

中与其同行同列的位置

。【

是向量的集合，在特征图

中各通道的位置

处通过交叉路径法获取的，也就是

中的一个位置

对应一个向量

，和一个向量集合

】因此，

。

是

的第

个元素【共有

个】。

关联操作定义如下：

——(1)

其中，

【实际是

，

】定义了特征

之间的相关性程度【

和交叉路径上各特征的相关程度】，

，

【取所有位置时，就变成了

】。然后，在通道维度上对

施加softmax层，以此计算注意力图

【尺寸为

，是基于

个通道上的各个位置得到的。】 。

然后，通过对

施加另一个卷积核为

的卷积层，以此生成用于特征适配的

。在特征图

的空间维度中的每个位置

，可以获取一个向量

和集合

【获取位置

的交叉路径上的特征向量得到的，其下标

共有

个】。集合

是特征图

中的特征向量集合，这些特征向量位于与位置

相同的行和列。远程上下文信息通过聚合操作获取：

——(2)

是一个标量，

为一个

的向量，

是交叉路径上向量的个数，

就是在各通道的交叉路径上的向量乘以对应的标量，结果的尺寸为

。而

的尺寸为

，所以可以相加。

其中，

定义了输出特征图

在位置

处的特征向量。

是注意力图

中通道

下位置

处的标量值。在局部特征

中添加了上下文信息，以此增强局部特征和增加像素级表示。因此， 局部特征

有广泛的上下文视野，依据空间注意力图选择性地聚合上下文信息。这种特征表示实现了互益性，对语义分割具有较强的鲁棒性。

交叉注意力模块是一个独立模块，可以在CNN结构的任何一个位置，嵌入任意数目的交叉注意力模块，以此获取丰富的上下文信息。该模块计算量非常低，并且仅添加了少量参数，因此显存占用率非常低。

方法-递归交叉注意力

尽管交叉注意力模块可以获取水平、垂直方向的远程上下文信息，像素与其邻近像素的连接仍然是稀疏的。获取密集上下文信息有益于语义分割。为了达到该目的，作者基于上述的交叉注意力模块，引入了递归交叉注意力。递归交叉注意力模块可以展开为一个

次循环。在第一次循环，交叉注意力模块将从CNN模型抽取出的特征图

作为输入，输出特征图

，

具有相同的形状。在第二次循环，交叉注意力模块将特征图

作为输入，输出特征图

。如图2所示，递归交叉注意力模块有两个循环(

)，足够获取所有像素的远程依赖，并生成新的密集的、上下文信息丰富的特征图。

是循环1、2中的注意力图。因为仅关注在空间维度上传播的上下文信息，而不是通道维度的，因此卷积核为

的卷积层可以视为等同连接。此外，位置

到权重

的映射函数定义为

。对于特征图

中任意位置

和特征图

中任意位置

，当

时仅有一个连接。一种情况是

和

在相同的行或列：

——(3)

是添加到操作。另一种情况是

和

在不同的行和列。图4给出了空间维度下的上下文信息传导路径：

——(4)

图 4 循环次数为2时，信息传导的例子

实际上，递归交叉注意力模块弥补了交叉注意力模块无法从所有像素中获取密集上下文信息的不足。与交叉注意力模块相比，递归交叉注意力模块(

)并未增加额外参数，并可以在增加少量计算量的代价下达到更好的性能。递归交叉注意力模块也是一个独立的模块，可以嵌入任意CNN结构的任意部分，并且可以通过端到端的方式优化。

实现细节

网络结构

作者基于开源的pytorch分割工具箱^[3]实现了其方法。为了实现语义分割，作者选择了基于ImageNet预训练的ResNet-101作为基干，移除了最后两个下采样操作，并在后续的卷积层中使用了扩张卷积，输出步长为8。同时，将标准的Batchnorm替换为InPlace-ABN^[4]，以此实现跨GPU计算Batchnorm的均值和标准差。为了实现实例分割，选择了Mask-RCNN^[5]作为基干。

训练设置

基于mini-batch 的 SGD 进行训练。对于语义分割，在 Cityscapes 和 ADE20K 上的初始学习速率为

。基于前人的工作

^{[6] [7]}，作者使用了多元学习速率策略，将初始学习速率乘以

。对 Cityscapes 中的训练图像进行了尺度为0.75至2.0的随机缩放，并从结果图像中裁剪出高分辨率块

。由于 ADE20K 中的图像具有多种尺寸，将输入图像的短边缩放为

中的一个随机值。此外，还使用了随机水平翻转。使用了4个 TITAN XP，以

进行训练。对于实例分割，使用同Mask-RCNN一致的训练方案。

参考

1. ^H. Zhao, Y. Zhang, S. Liu, J. Shi, C. C. Loy, D. Lin, and J. Jia. Psanet: Point-wise spatial attention network for scene parsing. In European Conference on Computer Vision, pages 270–286. Springer, 2018.
2. ^X.Wang, R. Girshick, A. Gupta, and K. He. Non-local neural networks. arXiv preprint arXiv:1711.07971, 10, 2017.
3. ^Z. Huang, Y. Wei, X. Wang, and W. Liu. A pytorch semantic segmentation toolbox. https://github.com/speedinghzl/ pytorch-segmentation-toolbox, 2018.
4. ^S. R. Bul`o, L. Porzi, and P. Kontschieder. In-place activated batchnorm for memory-optimized training of dnns. CoRR, abs/1712.02616, December, 5, 2017.
5. ^K. He, G. Gkioxari, P. Doll´ar, and R. Girshick. Mask r-cnn. In Computer Vision (ICCV), 2017 IEEE International Conference on, pages 2980–2988. IEEE, 2017.
6. ^L.-C. Chen, G. Papandreou, I. Kokkinos, K. Murphy, and A. L. Yuille. Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. IEEE transactions on pattern analysis and machine intelligence, 40(4):834–848, 2018.
7. ^H. Zhang, K. Dana, J. Shi, Z. Zhang, X.Wang, A. Tyagi, and A. Agrawal. Context encoding for semantic segmentation. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.