CCNet: Criss-Cross Attention for Semantic Segmentation论文阅读

CCNet: Criss-Cross Attention for Semantic Segmentation

Abstract

上下文信息对于语义分割和目标检测任务都很重要，这里提出CCNet。对于每个像素，criss-cross attention模块能获得其交叉路径上所有像素的上下文信息，通过进一步的递归操作，每个像素最终可以捕获全图像的依赖关系。此外，提出类别一致损失使得criss-cross attention模块生成更具判别性的特征。CCNet有以下优点：(1)GPU显存友好，比non-local block少11倍显存消耗 (2)高计算效率，比non-local block少85% (3)最先进性能，Cityscapes可达81.9%
开源地址：https://github.com/speedinghzl/CCNet
这里推荐一篇很有意思的工作，通过注意力机制处理预测head关联多目标跟踪领域的检测和ReID特征：Rethinking the competition between detection and ReID in Multi-Object Tracking，该论文不但构思了一种简单有效的注意力机制，还巧妙地利用注意力机制来交叉关联两个任务，避免了检测和ReID的竞争，并联合提升了彼此分支的性能。

Introduction

FCN是固定的几何结构(卷积的网格效应)，局部感受野只能提供短距离上下文信息。为了弥补FCN的缺点，deeplab系列提出多尺度空洞卷积结构的ASPP模块聚合上下文信息，PSPNet引入金字塔池化模块捕获上下文信息。然而，基于空洞卷积的方法从一些周围像素收集信息，不能准确产生稠密的上下文信息；基于池化的方法以非自适应的相同上下文提取策略处理所有像素，不能满足不同像素需要不同上下文依赖的要求。Non-local的空间和时间复杂度高，需要改进。

Related work

UNet，Deeplabv3+，MSCI，SPGNet，RefineNet，DFN采取encoder-decoder结构，融合低层次和高层次信息做出稠密预测。Scale-adaptive Convolutions(SAC)和Deformable Convolutional Network(DCN)改善标准卷积处理目标形变和各种尺寸目标…

Approach

1.Network Architecture

上图是Non-local block和Criss-Cross Attention block结构的简化示意图，Non-local block通过计算任意两个位置之间的交互直接捕获远距离依赖，而不用局限于相邻点，但是由于每个位置对应的向量(共$H\times W$个)都要和$H\times W$个向量相乘，带来的计算量偏大；而Criss-Cross Attention block一次只用考虑"十字交叉"的同行同列的向量，即每个位置对应的向量($H\times W$个)都要和$(H+W?1)$个向量相乘，这样捕获的是一个位置和”十字交叉“路径上其他位置的依赖，但是将输出结构果再次送入Criss-Cross Attention block即可获得一个位置与全局位置的依赖，进一步的理解可以看下文的Criss-cross Attention模块细节。

Backbone是全卷积网络，移去最后两个下采样操作，并且随后的卷积层都采取空洞卷积(带空洞卷积的FCN)，输出特征图$X$为输入图像的1/8。$X$通过卷积层降低通道维度输出$H$，$H$送入criss-cross attention模块聚合每个像素交叉路径上的上下文信息得到$H^{\prime }$，$H^{\prime }$再次送入criss-cross attention模块输出$H^{\prime \prime }$，则$H^{\prime \prime }$的每个像素聚合了所有像素的信息。两个criss-cross attention模块共享参数，取名为recurrent Criss-Cross Attention(RCCA)模块。然后，$H^{\prime \prime }$和特征$X$进行concat，接着是一个或几个带BN的卷积层和激活层用于特征融合，最后融合的特征送入分割层预测最终的分割结果。

2.Criss-cross Attention

考虑一个局部特征图$H\in {\mathbb{R}}^{C\times W\times H}$，首先通过两个$1\times 1$卷积生成两个特征图$Q$和$K$，{Q,K}∈RC′×W×H\{Q,K\} \in{\mathbb{R}^{C'\times W\times H}}{ Q,K}∈RC′×W×H，$C^{\prime }$是比$C$小的通道数，形状为$C^{\prime }\times H\times W$的三维特征图可以很容易reshape成二维的$C^{\prime }\times (H\times W)$的矩阵。通过Affinity操作生成注意力图$A$，对于特征图$Q$的每一个位置$u$，拉出一条维度为$C^{\prime }$的向量$Q_u\in {\mathbb{R}}^{C^{\prime }}$。同时，从特征图$K$拉出同属于$u$位置的同行或同列的$(H+W?1)$条($H+W$会包括$u$位置两次)维度均为$C’$的向量集${\Omega }_u\in {\mathbb{R}}^{(H+W?1)\times C^{\prime }}$，${\Omega }_{i,u}\in {\mathbb{R}}^{C^{\prime }}$是${\Omega }_u$的第$i$个元素(向量)，则Affinity操作可用公式表达如下：$$d_{i,u}=Q_u{\Omega }_{i,u}^T$$ $d_{i,u}\in D$是特征$Q_u$和${\Omega }_{i,u}$之间的关联度，$i=[1,...,H+W?1]$，且$D\in {\mathbb{R}}^{(H+W?1)\times (W\times H)}$，然后对$D$在通道维度上添加softmax层，输出注意力图A

输入特征图H通过另外一个$1\times 1$卷积生成特征图$V\in {\mathbb{R}}^{C\times W\times H}$，对于特征图$V$的每一个位置$u$，同理拉出维度为$C$的向量$V_u\in {\mathbb{R}}^C$和向量集 ${\Phi }_u\in {\mathbb{R}}^{(H+W?1)\times C}$，然后给出聚合(Aggregation)操作的公式如下：$$H_u^{\prime }=\sum _{i=1}^{H+W?1}{A}_{i,u}{\Phi }_{i,u}+H_u$$
其中，$H_u^{\prime }$是$H^{\prime }\in {\mathbb{R}}^{C\times W\times H}$中位置$u$的特征向量，$A_{i,u}$是注意力图A中位置$u$对应的第$i$个数值。最后是以残差的形式输出$H^{\prime }$，增强了像素级的表达能力，并聚合了全局上下文信息，提升了语义分割的性能。

Recurrent Criss-Cross Attention(RCCA)模块包含两个Criss-Cross Attention模块，且是共享参数的，RCCA可获得一个位置与全局位置的依赖，能够获得稠密丰富的下文信息。

3.Learning Category Consistent Features
对于语义分割任务，同一类像素应该有相似的特征，不同类像素应该有差别大的特征，这被称作类别一致性。论文认为，RCCA模块聚合的特征可能会存在过度平滑的问题，这是图神经网络的常见问题，因此除了使用交叉熵损失$l_{seg}$监督外，还提出了类别一致损失。RCCA模块输出后接$1\times 1$卷积降低特征图通道数(这里设置16)，在这个低通道数特征图$M$上添加类别一致性损失。假定$C$是mini-batch images里存在的类别数，$N_c$是属于类别$c\in C$的有效元素数目，$h_i\in H$是特征图M空间位置$i$对应的特征向量($i$是属于类别$c$的，是$N_c$中的一个元素)，$u_c$是类别$c$的平均特征向量(聚类中心)，${\phi }_{var}(h_i,u_c)$计算两者之间的距离进行惩罚，希望同类别像素对应的特征向量具有相似性，靠近该类聚类中心最好；$u_{c_a}$和$u_{c_b}$是两个不同类别的聚类中心（特征向量），${\phi }_{dis}(u_{c_a},u_{c_b})$计算两个类别中心之间的距离进行惩罚，两两类别计算，希望不同类别像素的聚类中心越远越好。$l_{reg}$是聚类中心向量的正则项损失，最终损失是所有损失的加权和：$$l=l_{seg}+\alpha l_{var}+\beta l_{dis}+\gamma l_{reg}$$其中，设置${\delta }_v=0.5,{\delta }_d=1.5,\alpha =\beta =1.\gamma =0.001$。总之，类别一致损失是从特征上，希望同类别像素特征具有相似性，不同类相似特征具有差异性。损失具体公式如下：