全图像依赖关系为视觉理解问题提供了有用的上下文信息。在这项工作中,我们提出了一种称为Criss-Cross Network(CCNet)的方法,以更有效和高效的方式获取这种上下文信息。具体而言,对于每个像素,CCNet中的一种新颖的交叉注意力模块收集其交叉路径上所有像素的上下文信息。通过进一步的循环操作,每个像素最终可以捕捉到来自所有像素的全图像依赖关系。总体而言,CCNet具有以下优点:
1)占用GPU内存较少。与非局部块相比,所提出的循环交叉注意力模块的GPU内存使用量减少了11倍。
2)计算效率高。循环交叉注意力在计算全图像依赖关系方面,将FLOPs的计算量显著减少了约85%。
3)达到了最先进的性能。我们在包括Cityscapes、ADE20K和COCO等流行的语义分割基准数据集上进行了大量实验。特别是,我们的CCNet在Cityscapes测试集和ADE20K验证集上分别取得了81.4和45.22的mIoU得分,这些都是最新的最优结果。
源代码可在https://github.com/speedinghzl/CCNet上获得。
语义分割是计算机视觉领域中的一个基本问题,其目标是为给定图像中的每个像素分配语义类别标签。它在许多最近的工作中得到了广泛而积极的研究,并且对于各种具有挑战性和有意义的应用,如自动驾驶[14]、增强现实[1]和图像编辑[13],也至关重要。具体而言,基于全卷积网络(FCN)[26]的当前最先进的语义分割方法取得了显著的进展。然而,由于固定的几何结构,它们天然受限于局部感受野和短程上下文信息。这些限制对基于FCN的方法造成了很大的不利影响,因为上下文信息不足。
为了弥补FCN的上述不足,一些工作已经提出引入有用的上下文信息来改善语义分割任务。具体而言,Chen等人[5]提出了带有多尺度膨胀卷积的空洞空间金字塔池化模块,用于上下文信息聚合。Zhao等人[41]进一步引入了带有金字塔池化模块的PSPNet来捕捉上下文信息。然而,基于膨胀卷积的方法[6,5,12]只从周围的几个像素收集信息,实际上无法生成密集的上下文信息。同时,基于池化的方法[41,39]以非自适应的方式聚合上下文信息,并且所有图像像素都采用相同的上下文信息,这不能满足不同像素需要不同上下文依赖的要求。
为了生成密集且像素级的上下文信息,PSANet [42]通过预测的注意力图来学习为每个位置聚合上下文信息。Non-local Networks [31]利用了自注意力机制[9,29],使得任何位置的单个特征可以感知到所有其他位置的特征,从而获取全图像的上下文信息,如图1(a)所示。然而,这些基于注意力的方法需要生成巨大的注意力图来度量每对像素之间的关系,其时间和空间复杂度都为O((H×W )×(H×W)),其中H×W表示输入特征图的空间维度。由于在语义分割任务中输入特征图始终具有高分辨率,基于自注意力的方法具有很高的计算复杂度并且占用大量的GPU内存。因此,是否有一种更高效的替代解决方案来实现这样的目标呢?
为了解决上述问题,我们的动机是通过连续的稀疏注意力替换非局部网络中的单层密集注意力。为了简单起见,我们使用了两个连续的交叉注意力模块,其中每个模块只对特征图中的每个位置进行(H+W-1)个稀疏连接。交叉注意力模块在水平和垂直方向上聚合上下文信息。通过串联两个交叉注意力模块,它可以从所有像素中收集上下文信息。上述分解策略将时间和空间复杂度从O((H×W)×(H×W))大大降低到O((H×W)×(H+W-1))。
我们在图1中比较了非局部模块[31]和我们的交叉注意力模块之间的区别。具体而言,非局部模块和交叉注意力模块都将具有空间尺寸H×W的输入特征图输入,分别生成注意力图(上分支)和适应性特征图(下分支)。然后,采用加权求和的方式收集上下文信息。与非局部模块采用的密集连接不同,我们的交叉注意力模块中的每个位置(例如蓝色)只与位于相同行和相同列的其他位置进行稀疏连接,导致预测的注意力图只有H+W-1个权重,而非局部模块中有H×W个权重。为了实现捕捉全图像依赖的目标,我们创新地对交叉注意力模块进行了简单的循环操作。具体而言,首先将局部特征通过一个交叉注意力模块传递,以在水平和垂直方向上收集上下文信息。然后,通过将第一个交叉注意力模块产生的特征图输入到另一个模块中,从交叉路径获得的额外上下文信息最终使得每个像素能够捕捉到全图像依赖关系。如图1(b)所示,第二个特征图中的每个位置(例如红色)最终从其他所有位置收集信息,以增强像素级表示。我们共享循环交叉模块的参数以减少额外的参数。我们的交叉注意力模块可以轻松地插入到任何完全卷积神经网络中,称为CCNet,以实现端到端的分割学习。
我们在多个大规模数据集上进行了大量实验。我们提出的CCNet在两个最具竞争力的语义分割数据集Cityscapes [10]和ADE20K [44]上取得了最佳性能。此外,我们提出的交叉注意力甚至改进了最先进的实例分割方法Mask R-CNN with ResNet-101 [17]。这些结果充分证明了我们的交叉注意力模块对于密集预测任务的普遍益处。总结起来,我们的主要贡献有两个方面:
图1.两种基于注意力的上下文聚合方法的示意图。
(a)对于每个位置(例如蓝色),非局部模块[31]生成一个密集的注意力图,其中有H×W个权重(绿色)。
(b)对于每个位置(例如蓝色),交叉注意力模块生成一个稀疏的注意力图,只有H+W-1个权重。经过循环操作后,最终输出特征图中的每个位置(例如红色)都可以从所有像素中收集信息。为了清晰显示,忽略了残差连接。
语义分割:语义分割近年来重新引起了人们的兴趣。FCN [26]是第一个采用完全卷积网络进行语义分割的方法。后来,基于FCN的方法在图像语义分割方面取得了巨大进展。Chen等人[4]和Yu等人[37]移除了最后两个下采样层以获得密集预测,并利用扩张卷积来扩大感受野。Unet [28]、Deeplabv3+ [8]、MSCI [21]、SPGNet [2]、RefineNet [22]和DFN [36]采用了编码器-解码器结构,融合低层和高层的信息来预测分割掩码。SAC [40]和可变形卷积网络 [11]改进了标准卷积算子,以处理物体的变形和不同尺度。CRF-RNN [37]和DPN [25]使用图模型,即CRF、MRF,进行语义分割。AAF [19]使用对抗学习来捕捉和匹配标签空间中相邻像素之间的语义关系。BiSeNet [35]是为实时语义分割而设计的。
上下文信息聚合:此外,一些工作通过聚合上下文信息来增强特征表示。Deeplabv2 [5]提出了ASPP模块,使用不同的扩张卷积来捕捉上下文信息。DenseASPP [34]将密集连接引入ASPP中,生成具有不同尺度的特征。DPC [3]利用架构搜索技术构建了用于语义分割的多尺度架构。PSPNet [41]利用金字塔池化来聚合上下文信息。GCN [27]利用全局卷积模块和全局池化来收集上下文信息用于全局表示。最近,Zhao等人[42]提出了逐点空间注意力网络,使用预测的注意力图来引导上下文信息的收集。Liu等人[24]利用循环神经网络来捕捉长距离的依赖关系。条件随机场(CRF)[4,43]、马尔可夫随机场(MRF)[25]也被用来捕捉语义分割中的长距离依赖关系。
注意力模型:注意力模型在各种任务中得到了广泛应用。Squeeze-and-Excitation Networks [18]通过建模通道间关系的注意力机制增强了网络的表示能力。Chen等人[7]利用多个注意力掩码来融合不同分支的特征图或预测结果。Vaswani等人[29]在机器翻译中应用了自注意力模型。Wang等人[31]提出了非局部模块,通过计算特征图中每个空间点之间的相关矩阵来生成大规模的注意力图,然后通过注意力引导的密集上下文信息聚合。OCNet [38]和DANet [15]利用非局部模块 [31]收集上下文信息。PSA [42]学习了一个注意力图,以自适应和具体地聚合每个个体点的上下文信息。
CCNet vs.Non-Local vs.GCN:在这里,我们特别讨论了GCN [27]、非局部网络 [31]和CCNet之间的区别。在上下文信息聚合方面,只有GCN [27]中的中心点可以感知到所有像素的上下文信息。相比之下,非局部网络 [31]和CCNet确保任何位置的像素都可以从所有像素中感知上下文信息。尽管GCN [27]将方形卷积操作分解为水平和垂直线性卷积操作,这与CCNet有关,但CCNet采用了交叉方式来获取上下文信息,比水平-垂直分离方式更有效。此外,CCNet旨在模拟非局部网络 [31],通过更有效和高效的循环交叉注意力模块来获得密集的上下文信息,其中不相似的特征获得低的注意力权重,而具有高注意力权重的特征是相似的。
图2.提出的用于语义分割的CCNet概览
在本节中,我们详细介绍了用于语义分割的Criss Cross网络(CCNet)的细节。首先,我们提出了CCNet的一般框架。然后,介绍了横向和纵向方向上捕捉上下文信息的交叉注意力模块。最后,为了捕捉密集和全局的上下文信息,我们提出采用循环操作来实现交叉注意力模块
网络架构如图2所示。输入图像通过一个深度卷积神经网络(DCNN)进行处理,该网络采用全卷积的方式设计[5],以产生具有H×W空间尺寸的特征图X。为了保留更多细节并有效地生成密集的特征图,我们去掉了最后两个下采样操作,并在后续的卷积层中使用了扩张卷积,从而将输出特征图X的宽度/高度扩大到输入图像的1/8。
给定特征图X,我们首先应用一个卷积层来获取维度减小的特征图H,然后将特征图H输入到交叉注意力模块中,生成聚合了每个像素在其交叉路径上的上下文信息的新特征图H'。特征图H'只聚合了水平和垂直方向上的上下文信息,这对于语义分割来说并不足够强大。为了获取更丰富和更密集的上下文信息,我们将特征图H'再次输入到交叉注意力模块中,输出特征图H''。因此,特征图H''中的每个位置实际上都汇集了来自所有像素的信息。前后两个交叉注意力模块共享相同的参数,以避免添加过多的额外参数。我们将这个循环结构命名为循环交叉注意力(RCCA)模块。
接下来,我们将密集的上下文特征H''与局部表示特征X进行拼接。然后,通过一个或多个带有批归一化和激活函数的卷积层进行特征融合。最后,融合后的特征被输入到分割层中,以预测最终的分割结果。
图3.交叉注意力模块的详细信息。
为了利用轻量级的计算和内存来建模局部特征表示上的全图依赖关系,我们引入了一个交叉注意力模块。交叉注意力模块在水平和垂直方向上收集上下文信息,以增强像素级别的代表能力。如图3所示,给定一个局部特征图H ∈RC×W ×H,该模块首先在H上应用两个1×1的卷积层,分别生成两个特征图Q和K,其中{Q,K}∈RC×W ×H。C是通道数,用于降维,比C小。
在获取特征图Q和K之后,我们通过亲和力操作进一步生成注意力图A ∈R(H+W −1)×W ×H。在特征图Q的空间维度的每个位置u上,我们可以获得一个向量Qu ∈RC。同时,我们还可以通过从K中提取与位置u在同一行或同一列的特征向量,得到集合Ωu ∈R(H+W −1)×C。Ωi,u ∈RC是Ωu的第i个元素。亲和力操作定义如下:
其中,di,u ∈D表示特征Qu和Ωi,u之间的相关度,i =[1,...,|Ωu|],D ∈R(H+W −1)×W ×H。然后,我们在di的通道维度上应用softmax层,计算注意力图A。
在H上应用另一个1×1的卷积层,生成V ∈RC×W ×H,用于特征适应。在特征图V的空间维度的每个位置u上,我们可以获得一个向量Vu ∈RC和一个集合Φu ∈R(H+W −1)×C。集合Φu是V中与位置u在同一行或同一列的特征向量的集合。通过聚合操作收集上下文信息:其中,Hu表示在位置u处的输出特征图H∈RC×W ×H中的一个特征向量。Ai,u是在A中的通道i和位置u处的标量值。上下文信息被添加到局部特征H中,以增强局部特征并增加像素级表示。因此,它具有广泛的上下文视角,并根据空间注意力图选择性地聚合上下文。这些特征表示实现了相互增益,并对语义分割更具鲁棒性。
图4. 循环次数为2时信息传播的示例
尽管交叉注意力可以在水平和垂直方向上捕捉上下文信息,但一个像素与不在交叉路径上的周围像素之间的连接仍然缺失。为了解决这个问题,我们创新地并简单地引入了基于交叉注意力的递归交叉注意力(Recurrent Criss-Cross Attention,RCCA)操作。RCCA模块可以展开成R个循环。在第一个循环中,交叉注意力以从CNN模型提取的特征图H作为输入,并输出特征图H,其中H和H具有相同的形状。在第二个循环中,交叉注意力以特征图H作为输入,并输出特征图H。如图2所示,RCCA模块配备了两个循环(R=2),能够从所有像素中收集全图上下文信息,生成具有密集且丰富上下文信息的新特征图。
我们将循环1和循环2中的注意力图表示为A和A。由于我们只关注空间维度上的上下文信息传播,而不关注通道维度,1×1卷积层可以视为相同的连接。此外,从位置x、y到权重Ai,x,y的映射函数被定义为Ai,x,y = f(A,x,y,x,y)。对于特征图H中的任意位置u和特征图H中的任意位置θ,在R = 2的情况下实际上存在一个连接。对于u和θ在同一行或同一列的情况:其中←表示加法操作。对于u和θ不在同一行和同一列的情况,图4显示了空间维度上上下文信息的传播路径。
总体而言,我们的递归交叉注意力(RCCA)模块弥补了交叉注意力无法从所有像素中获取密集上下文信息的不足。与交叉注意力相比,RCCA模块(R = 2)不会增加额外的参数,并且可以在稍微增加计算量的情况下实现更好的性能。
为了评估CCNet的有效性,我们在Cityscapes数据集[10]、ADE20K数据集[44]和COCO数据集[23]上进行了全面的实验。实验结果表明,CCNet在Cityscapes和ADE20K数据集上实现了最先进的性能。同时,CCNet在COCO数据集上对实例分割也能带来持续的性能提升。在下面的小节中,我们首先介绍数据集和实现细节,然后在Cityscapes数据集上进行一系列消融实验。最后,我们报告了在ADE20K和COCO数据集上的结果。
我们采用平均交并比(mIOU,类间交并比的平均值)作为Cityscapes和ADE20K的评估指标,采用标准的COCO指标平均精确度(AP)作为COCO的评估指标。
网络结构 对于语义分割任务,我们选择使用ImageNet预训练的ResNet-101[17]作为我们的主干网络,并移除最后两个下采样操作,在随后的卷积层中采用扩张卷积(dilated convolutions),这是根据之前的工作[4]进行的改进,使输出步幅为8。 对于实例分割任务,我们选择使用Mask RCNN[16]作为我们的基线模型。
训练设置 我们使用SGD(随机梯度下降)和小批量训练的方法进行训练。 对于语义分割任务,Cityscapes和ADE20K的初始学习率为1e-2。我们采用了先前的工作[5,39]中使用的多项式学习率策略,初始学习率乘以(1 - max_iter/iter)^power,其中power=0.9。我们使用动量为0.9,权重衰减为0.0001。 对于Cityscapes数据集,训练图像通过随机缩放(从0.75到2.0),然后从结果图像中随机裁剪出高分辨率补丁(769×769)来增强数据。由于ADE20K数据集中的图像大小各异,我们采用了一种缩放策略,将输入图像的短边缩放到从集合{300,375,450,525,600}中随机选择的长度。 对于实例分割任务,我们采用了与Mask-RCNN[16]相同的训练设置。
其他最先进的语义分割方法在Cityscapes验证集上的结果总结如表1所示。我们提供这些结果供参考,并强调不应简单地将这些结果与我们的方法进行比较,因为这些方法是在不同(甚至更大)的训练集或不同的基础网络上训练的。在这些方法中,Deeplabv3[6]和CCNet是其中性能最好的两种方法。
此外,我们还使用ResNet-101作为主干网络对最佳学习的CCNet进行了训练,同时使用训练集和验证集进行了评估,并通过将测试结果提交给官方评估服务器来对测试集进行评估。大多数方法[5,22,40,27,30,41,35,19,42,36]采用与我们相同的主干网络,其他方法[32,34]则使用更强大的主干网络。从表2可以看出,我们的CCNet在性能上大大优于之前的所有最先进方法。在这些方法中,PSANet[42]与我们的方法最相关,它为每个像素生成一个子注意力图。其中一个区别是,PSANet的子注意力图具有2×H×W个权重,而CCNet的子注意力图具有H+W-1个权重。即使在计算成本和内存使用方面更低的情况下,我们的方法仍然能够实现更好的性能。
使用train-fine和val-fine数据集进行训练。
表3.在Cityscapes(验证集)上对不同循环次数的RCCA性能进行评估。FLOPs和内存增加量是根据输入为1×3×769×769进行估计的。
为了验证CCNet的合理性,我们在Cityscapes的验证集上进行了广泛的削减实验,使用不同的CCNet设置。 我们通过对CCNet进行不同设置的实验来验证其合理性。我们在Cityscapes的验证集上进行了这些实验。通过对CCNet进行削减,我们可以评估每个设置对性能的影响。这些实验可以帮助我们了解CCNet的不同组件和参数对结果的影响,并找到最佳的设置。通过这些削减实验,我们可以验证CCNet的设计和性能,并进一步优化和改进它。
RCCA模块的效果表3显示了在Cityscapes验证集上采用不同循环次数的RCCA的性能。所有实验都是使用ResNet-101作为骨干网络进行的。此外,每个图像的输入尺寸为769×769,导致RCCA的输入特征图H的尺寸为97×97。我们的基准网络是基于ResNet的FCN,其中在第4和第5阶段分别引入了扩张卷积模块,即这两个阶段的扩张率分别设置为2和4。当R = 1、2、3时,估计了FLOPs和内存使用的增加量。
我们观察到,将交叉注意力添加到基准模型中(R=1)相对于基准模型可以提高2.9%的性能,这有效地证明了交叉注意力的重要性。此外,将循环次数从1增加到2可以提高1.8%的性能,说明了密集上下文信息的有效性。最后,将循环次数从2增加到3会略微提高0.4%的性能。同时,随着循环次数的增加,FLOPs和GPU内存的使用也会增加。这些结果证明了提出的交叉注意力能够通过捕捉水平和垂直方向上的上下文信息显著提高性能。此外,提出的交叉注意力在捕捉密集和全局上下文信息方面是有效的,最终有助于语义分割的性能。为了平衡性能和资源使用,我们在所有后续实验中选择R=2作为默认设置。
为了进一步验证交叉注意力模块的有效性,我们在图5中提供了定性比较。我们使用白色圆圈来指示那些容易被错误分类的具有挑战性的区域。可以看到,随着循环次数的增加,这些具有挑战性的区域逐渐得到修正,这充分证明了密集上下文信息聚合对语义分割的有效性。
不同的上下文聚合方法:我们在Cityscapes验证集上使用ResNet-50和ResNet-101作为骨干网络,比较了几种不同的上下文聚合方法的性能。具体来说,上下文聚合的基线方法主要包括:
1)Peng等人[27]使用全局卷积网络进行上下文信息聚合,标记为“+GCN”;
2)Zhao等人[41]提出了金字塔池化方法,这是一种简单而有效的捕捉全局上下文信息的方法,标记为“+PP”;
3)Chen等人[6]使用不同的扩张卷积在不同范围内收集像素级别的上下文信息,标记为“+ASPP”;
4)Wang等人[31]引入了非局部网络进行上下文聚合,标记为“+NL”。
在表4中,与其他上下文聚合方法相比,"+NL"和"+RCCA"都取得了更好的性能,这证明了捕捉全图上下文信息的重要性。更有趣的是,我们的方法比"+NL"表现更好。这可能归因于交叉注意力模块的顺序循环操作。具体来说,"+NL"直接从具有有限感受野和短程依赖性的特征中生成注意力图。相反,我们的"+RCCA"需要两个步骤来形成密集的上下文信息,导致后一步可以从第一步产生的特征图中学习到更好的注意力图,其中已经嵌入了一些长程依赖性。
为了证明交叉形状的注意力的有效性,我们在表4中将交叉形状与其他形状进行了比较。"+HV"表示堆叠水平注意力和垂直注意力。"+HV&VH"表示将两个并行分支"HV"和"VH"的特征相加。这些结果证明了交叉注意力可以比其他形状取得更好的性能。我们进一步探索了RCCA的计算量和内存占用量。如表5所示,与"+NL"方法相比,提出的"+RCCA"方法在计算全图依赖性时需要11倍更少的GPU内存使用,并且FLOPs显著减少了约85%。这表明CCNet是一种在最少的计算量和内存占用量下捕捉全图上下文信息的高效方法。
可视化注意力图:为了更深入地了解我们的RCCA,我们在图6中可视化了学习到的注意力掩码。对于每个输入图像,我们选择一个点(绿色的十字)并在第二列和第三列分别显示其在R=1和R=2时的相应注意力图。可以观察到,当R=1时,只有来自目标点交叉路径的上下文信息被捕捉到。通过采用一个更多的交叉注意力模块,即R=2,RCCA最终可以聚集更密集、更丰富的上下文信息,相比于R=1的情况。此外,我们观察到注意力模块能够捕捉语义相似性和全图依赖性。
图5.在Cityscapes验证集上使用不同循环的RCCA的可视化结果。
表4.在Cityscapes(验证集)上对上下文聚合方法的比较。
图6.在Cityscapes验证集上对注意力模块的可视化。左列是输入图像,第2列和第3列是在RCCA中R=1和R=2时的像素级注意力图。表5.非局部模块和RCCA的比较。FLOPs和内存增量是根据输入
表5.非局部模块和RCCA的比较。FLOPs和内存增量是根据输入为1×3×769×769进行估计的。为1×3×769×769进行估计的。
在本小节中,我们在AED20K数据集上进行了实验,该数据集是一个非常具有挑战性的场景解析数据集。如表6所示,CCNet实现了45.22%的最新性能,优于之前最先进的方法超过0.6%。在这些方法中,大多数方法[40,41,42,20,33,39]采用ResNet-101作为主干网络,而RefineNet [22]采用了更强大的ResNet-152作为主干网络。EncNet [39]在这些方法中实现了最好的性能,并利用全局池化和图像级监督来收集图像级的上下文信息。相比之下,我们的CCNet采用了一种替代的方式来整合上下文信息,通过捕捉全图像依赖关系来实现更好的性能。
为了进一步证明CCNet的普适性,我们使用竞争的Mask R-CNN模型[16]作为基准,在COCO [23]上进行了实例分割任务。根据[31]的方法,在res4的最后一个卷积残差块之前添加了RCCA模块来修改Mask R-CNN的主干网络。我们评估了ResNet-50/101的标准基准模型。所有模型都是从ImageNet预训练中微调得到的。我们使用了官方实现1进行端到端联合训练,其性能与[31]中报告的基准模型几乎相同。我们在COCO上根据盒子AP和掩码AP报告了结果,如表7所示。结果表明,我们的方法在所有指标上显著优于基线模型。同时,带有“+RCCA”的网络也比带有一个非局部块“+NL”的网络表现更好。
表7.在COCO(验证集)上的比较
在本文中,我们提出了一种用于基于深度学习的密集预测任务的Criss-Cross网络(CCNet),它能够自适应地在交叉路径上捕获上下文信息。为了获得密集的上下文信息,我们引入了RCCA,它从所有像素中聚合上下文信息。实验表明,RCCA以较低的计算成本和内存成本捕捉全图像的上下文信息。我们的CCNet在两个语义分割数据集(Cityscapes、ADE20K)和一个实例分割数据集(COCO)上始终取得出色的性能。