横着注意力+竖着注意力

作者:Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille,and Liang-Chieh Chen    

单位:约翰霍普金斯大学,谷歌                                        ECCV2020 spotlight

https://arxiv.org/pdf/2003.07853v2.pdf

https://github.com/lucidrains/axial-attention

https://github.com/csrhddlam/axial-deeplab

摘要:卷积以丢失远程上下文为代价,利用局部性来提高效率。自注意力已被用来作为卷积的补充来捕获远程关系。最近的工作表面,通过将注意力限制在局部区域,可以堆叠注意力层来实现一个全注意力网络。本文通过将二维注意力分解为2个一维注意力来消除这种限制,这样降低了计算量,并且可以执行更大区域甚至全局区域的注意力。我们还设计了一个位置敏感的自注意力模块。将这两种方法结合起来,就是我们的位置敏感的轴向注意力层。这是种全新的构建块,可堆叠以形成用于图像识别及密集预测(分割、检测等)的轴向注意力模型。在4个大规模数据集上证明了有效性:在ImageNet上好于所有的独立的自注意力模型;Axial DeepLab在COCO test-dev,比自下向上的先前最先进模型提高了2.8%的PQ。这一先前的最先进模型通过我们的微小改变,参数量缩小3.8倍、计算量缩小27倍。Axial DeepLab在Mapillary Vistas和Cityscapes上也达到最先进水平。

关键词:自底向上全局分割,自注意力

1 引言

卷积是计算机视觉的核心组成部分。传统图像算法使用卷积滤波器来模糊图像、提取边缘或检测特征。与全连接网络相比,卷积网络因其计算量小、泛化能力好而在现代神经网络[50,49]中得到广泛使用[2]。卷积的成功主要来自两个性质:平移的等变性和局部性。平移等变性,虽然从严格意义上是不精确的[97],但与成像的性质很好地吻合,因此模型能够适用于图像的不同位置,能够适用于不同尺寸的图像。另一方面,局部性减少了参数量和计算量,但同时也使得远程关系建模更困难。

有很多工作是研究在卷积网络中远程关系建模。有些使用空洞卷[34,77,67,13],大卷积核[70],以及人工设计或搜索[103,12,60]得到的图像金字塔[98,85]。另一类方法是使用注意力机制。注意力在语言建模[83,88]、语音识别[22,11]和神经字幕[92]中显示了其远程关系建模的能力。之后,注意力被运用到视觉领域,对图像分类[6]、目标检测[37]、语义分割[42]、视频分类[87]和对抗性防御[89]都带来很大的推动。这些工作使用non-local或者说远程注意力模块,大大丰富了卷积网络。

最近出现了堆叠注意力层的网络,是一个无需任何空间卷积的独立(stand-alone)注意力模型,并且效果不错。然而,原始的注意力模块的计算量是很大的,尤其当输入的尺寸很大的时候。[68,38]对注意力增加了局部约束,也就是局部注意力,可降低计算量,并且建立全注意力模型。然而,局部约束导致模型感受野小,而大的感受野对分割任务及大尺寸输入的情况是至关重要的。本文提出使用轴向注意力[33,42],不仅减小计算量,而且可在全注意力网络中恢复大感受野。核心思想是将二维注意力依次沿着竖直方向和水平方向分解为两个一维注意力。分解使得计算量小,这样我们就可以让注意力覆盖到大的区域,并且建模远程甚至全局的关系。除此之外,先前的注意力模块都没有利用位置信息,这样就降低了注意力模块在与位置相关的关系中的建模能看,例如形状,以及出现不同大小的物体。近期的工作[68,38,6]再注意力模块中增加位置项(考虑了位置信息),但是是以一种上下文无关的方式。本文我们增加了一种依赖于上下文的位置项,使得我们的注意力模块在只增加少量计算量的情况下对位置敏感。

我们在ImageNet[73]分类以及三个数据集(COCO[59]、Mapillary Vistas[65]和Cityscapes[23])的分割上展示了我们的方法是有效的。特别的,在ImageNet分类上,我们建立了Axial-ResNet,是将所有的残差块[32]中的3×3卷积替换为我们的位置敏感轴向注意力块。我们更进一步地,在网络的“柄”上使用我们的轴向注意力层,将网络变成一个全注意力网络。结果表明,我们的Axial-ResNet在ImageNet的所有全注意力网络中取得最先进水平。在分割任务上,我们通过将主干网络改为Panoptic-DeepLab,将Axial-ResNet改为Axial-DeepLab。在COCO数据集上, Axial-DeepLab要比当前的自底向上最先进方法Panoptic-DeepLab[20]更好,在test-dev set上要高出2.8%PQ。我们还在Mapillary Vistas [65]和Cityscapes [23]上取得最先进水平。

总之,我们的贡献有四个方面:

——首次建立具有大感受野和全局感受野的全注意力;

——提出位置敏感的注意力层,可以更好地利用位置信息而没有带来太多的计算量;

——轴向注意力不仅可以用于图像分类,还可作为全局分割、实例分割和语义分割的主干;

——在COCO上比当前的自底向上最先进方法有显著提升,与两阶段方法匹敌;在Mapillary Vistas和Cityscapes上也超越了之前的最先进方法。

2 相关工作

自顶向下的全景分割:    两阶段方法,先输出object proposals,再依次按顺序处理。略

自底向上的全景分割:    无object proposals,先语义分割,接着将“thing”像素分组成簇以获得实例分割。DeeperLab[93]预测了边界框的四个角和中心,用于类别无关的实例分割。SSAP[29]利用了像素对相似金字塔[63],该金字塔由一种高效的图(graph)划分方法[46]实现。BBFNet[8]通过分水岭变换[84,4]和Hough投票机制[5,51]得到实例分割结果。最近,Panoptic DeepLab[20]是一种简单、快速、强大的自底向上全景分割方法,它采用了一个与类别无关的实例分割分支,包含一个简单的实例中心回归[45,82,66],再加上DeepLab语义分割输出[13,15,16]。Panoptic DeepLab已经在多个基准测试中取得了最先进的结果,我们的方法建立在它之上。

自注意力:    注意力,由[3]提出是在序列建模中用于编码解码的,用于捕获序列中两个位置的关系。与之不同的是,自注意力是在一个context上使用注意力(也就是key和query来自同一个feature),而不是跨越多个模式。它直接建模远程关系,并且计算并行化,因而在许多任务上达到最先进水平[83,40,26,69,75,25,56]。近年来,自注意力被用于计算机视觉,通过增加non-local或者远程建模来补充卷积网络。non-local网络[87]表明,自注意力是non-local的一个实例,在许多任务例如视频分类和物体检测上都能带来提升。另外,[18,6]通过将卷积输出的特征和自注意力输出的特征组合起来,在图像分类上得到了提升。最先进的视频动作识别模型[18]也使用了这样的方法。在语义分割领域,自注意力是作为一个上下文聚合模块,用于捕获多尺度的上下文[42,27,102,99]。[76,42,56]提出了计算量小的注意力模块。另外,带有non-local的卷积网络被证明在对抗攻击下是更鲁棒的[89]。除了识别任务,自注意力也用于图像生成[95,9,33]。最近[68,38]表明,将自我注意的感受野限制在一个局部的正方形区域,仅仅使用自注意力层的堆叠可以构建一个全注意力模型,在图像分类和物体检测上都取得鼓舞人心的结果。本文是沿着这一研究方向,提出一个具有大感受野和全局感受野的自注意力模型,使自注意力再一次non local。我们的模型在自底向上全景分割中评估,并显示出显著的改进。

3 方法

先介绍位置敏感的自注意力机制,再讨论如何将其用到轴向注意力上,最后在使用轴向注意力层建立Axial-ResNet和Axial-DeepLab。

3.1 位置敏感的自注意力

自注意力:    自注意力常被用到视觉模型中作为一个可相加模块,用于补充卷积网络的输出[87,95,42]。给定一个输入特征图,形状为,也就是高为,宽为,通道数为。在位置处的输出,形状为,是将输入汇集为:

    

其中是所有位置,查询项、键项、值项都是输入的线性变换。和的形状为,的形状为,它们都是可学习的参数矩阵。表示在所有可能的位置上使用softmax函数,在本文的情况下也就是空间维度上所有位置。

这个机制将值项根据亲缘关系全局地汇聚起来,使得我们可以在整个特征图上捕获non-local的上下文关系,这和卷积不同,卷积只能捕获局部的上下文关系。

然而,当输入尺寸大的时候,自注意力的计算复杂度会很大,这使得它只被用于卷积网络的高级特征(也就是下采样过的特征)或者小图像上。另一个缺点是,全局汇聚不能利用位置信息,而在要求捕获空间结构和形状的视觉任务中,位置信息是关键的。

这两个问题在[68]中通过增加局部约束和位置编码来减轻。对于每个位置,只计算周边的区域,这样计算复杂度就大大减小至,这样就使得建立全注意力网络可行。此外,将一个可学习的相对位置编码项包含到亲缘关系中,这样就产生中一个在感受野(也就是区域)中应该朝哪儿看的动态先验。[68]提出的正式计算形式是:

    

其中,就是以位置为中心的大小为的局部正方形区域。可学习的向量就是附加的相对位置编码器。内积衡量的是从位置到位置的相容性。我们没有考虑绝对位置编码,因为其泛化能力不如相对位置编码。在下文中,为了行文简洁,我们省略了“相对”这个词。

在实际中,和要比要小得多,因此很容易将公式的单头注意力扩展为多头注意力,以捕获多种亲缘关系的混合。具体的实现是,多头注意力对输入并行地计算个单头注意力,也即,然后将每个单头注意力的输出拼接起来得到最终的输出。注意,不同的注意力头常常用的是同一个位置编码,这用引入的额外计算量就很小。

位置敏感:    我们注意到先前的位置偏置主要是依赖于查询像素而不是键像素。然后,键像素也能够提供应该关注哪里的信息。因此,我们除了依赖于查询项的位置偏置,还增加一项依赖于键的位置偏置。

类似地,公式(2)中的值项没有包含任何的位置信息,在大感受野的情况下,不太可能包含精确的从哪里来的位置信息。因此,之前的方法需要在小的感受野(也就是小的区域)和精细的空间结构的损失之间做一个权衡。在本文中,我们使得输出除了能利用,还能够检索相对位置,基于查询-键的亲缘关系上。计算公式是:

   

其中,可学习的是键的位置编码,是值的位置编码。两个向量都没有引入太多参数,因为它们在每层中是共享同一个注意力头,而局部像素通常都很少。

我们将其称为位置敏感的自注意力,在我们的实验验证中,它能捕获远程关系,能包含精确的位置关系,计算量在合理范围内。

图1 右边是我们的沿着宽度方向的位置感知的轴向注意力,左边是一个non local块。 表示矩阵相乘, 表示元素按位相加。softmax的计算是在最后个axis。蓝框表示1x1卷积,红框表示相对位置编码。图中的通道数d_in = 128,d_q = 8, d_out = 16,是我们在ResNet卷积柄之后的第一个阶段的设置

3.2 轴向注意力

[68]提出局部约束大大降低了在视觉任务上的计算量,使得建立全注意力网络可行。但是这样的约束牺牲了全局关系,使得注意力的感受野不会大于相同核大小的深度卷积。此外,局部自注意力是在正方形区域执行注意力,其计算复杂度正比于正方形边长,因此又带来另一个超参数来平衡性能和计算量。本文提出在全注意力层中的自注意力中采用轴向注意力,以保证全局的关系建模和低计算量。具体的是,在图像的宽的方向上定义轴向注意力层为一维的位置敏感自注意力,然后在图像的高的方向上做类似的定义。具体的是,在宽的方向上的轴向注意力定义为:

    

一个轴向注意力层沿着一个特定的轴的方向传播信息。为了捕获全局信息,先是在宽度方向上,再在高度方向上接连地使用两个轴向注意力层。如上文所述,两个轴向注意力层使用多头注意力机制。

轴向注意力将计算复杂度减小至。通过将跨度直接设置为整个输入特征的宽(或高),就可以实现全局感受野。或者,也可以将设为固定值,以减少在大的特征图上的内存占用。

图2 一个轴向注意力块,其包含2个轴向注意力层,分别计算高度方向和宽度方向。通道数d_in = 128,d_out = 16,是我们在ResNet的卷积柄后第一个阶段的设置。我们的注意力头数设置为N = 8

Axial-ResNet:    为了将ResNet[32]变成Axial-ResNet,我们将residual bottleneck block中的3×3卷积替换为2个多头轴向注意力层,一个用于高度方向,一个用于宽度方向。在每个轴方向上,还可使用striding(选择性使用),在相应的轴向注意力层之后。两个1×1卷积保留,用以对特征洗牌。这样就形成了我们的(残差)轴向注意力块,如图2所示。然后将这个块堆叠多次,就形成了Axial-ResNets。注意,我们没有在两个轴向注意力层之间使用1×1卷积,因为矩阵乘法是紧接着的。此外,原始的ResNet的柄(即前2层步幅为2的7×7卷积和3×3最大池化)会被保留,这样就形成了一个卷积柄网络(conv-stem),也就是卷积只在第一层中使用,注意力则在其余的各个层使用。在卷积柄网络中,我们将跨度设置为第一个块的输入宽/高,在这里特征图的尺寸是56×56。

在我们的实验中,我们还建立了一个全-轴向注意力网络,称为Full-Axial-ResNets,在柄中也使用轴向注意力。不同于[68]中专门设计一个在空间维度上变化的注意力柄,我们只是简单地将3个轴向注意力瓶颈块堆叠起来,作为注意力柄。此外,我们还在Full Axial-ResNets的前几个块中还使用了局部约束,也就是[68]中的m×m的正方形领域中计算注意力,以减小计算量。

Axial-DeepLab:    对于分割任务,将Axial-ResNet改成Axial-DeepLab,我们进一步地做以下更改。

首先,为了提取密集的特征图,DeepLab[13]改变了ResNet[32]中最后一个或两个阶段的步幅和扩张率。同样,我们去掉了最后一个阶段的步幅(也就是将步幅从2改为1),但我们没有实现扩张注意力模块,因为我们的轴向注意力已经捕捉到全局的信息。我们将输出步幅(即输入的全分辨率与主干网络最终输出的特征图的分辨率之比)设为16,而不是8,因为设为8的话计算量大。(笔者:注意CItyscapes等的输入分辨率是很大的。)

其次,我们没有实现扩张空间金字塔池化模块(ASPP)[14,15],因为我们的轴向注意力模块也能高效地编码多尺度及全局的信息。实验表明,没有ASPP的Axial-DeepLab,比使用或不使用ASPP的Panoptic-DeepLab[20]的性能都要好。

最后,我们使用了和Panoptic-DeepLab[20]同样的stem[81]:3个卷积、双编码器、预测头。预测头输出语义分割和类别无关的实例分割,并按照多数投票[93]进行合并,形成最终的全景分割。

在输入非常大(例如2177x2177)和内存受限的情况下,我们在所有的轴向注意力块中使用大跨距m=65。因为它在几个数据集上足以覆盖远程甚至全局的上下文,而且,将m取小的值也不会显著减小M-Adds。

4 实验结果

我们在四个大型数据集上进行了实验。我们首先在ImageNet上使用轴向ResNet报告结果[73]。然后,我们将ImageNet预处理的轴向ResNet转换为Axial DeepLab,并通过全景质量(PQ)[48]评估COCO[59]、Mapillary Vistas[65]和Cityscapes[23]的结果。我们还报告了实例分割的平均精度(AP),以及Mapillary Vistas和Cityscapes语义分割的mIoU。我们的模型使用TensorFlow[1],在ImageNet上使用128TPU核,在全景分割上使用32核。

训练流程超参:    在ImageNet上,为了公平比较,我们采用了与[68]相同的训练流程与超参,有一点不同的是,我们在Full Axial-ResNets上将批量大小设为512,而其它模型的批量大小设为1024,学习率做相应的调整[30]。

在全景分割上,我们严格遵循Panoptic-DeepLab[20]的超参,不同之处是我们使用线性warm-up Radam[61] Lookhead[96]优化器,学习率和[20]一样设为0.001。我们的所有全景分割模型都使用这个设置。我们注意到,对优化器的改动其实并没有提升结果,但是平滑了训练曲线。Panoptic DeepLab在这种设置下有出现类似结果。

4.1 ImageNet

在ImageNet上的实验,我们从ResNet-50[32]建立Axial-ResNet-L网络。具体地是,在柄之后的第一个阶段,我们设置,。当特征分辨率宽高减半时,将这些取值翻倍。另外,我们将各通道数缩放0.5、0.75、2,就产生了不同大小版本的Axial-ResNet-{S,M,XL}[36,74,35]。最后,将卷积柄换成3个轴向注意力块,就进一步地建立了Stand-Alone Axial-ResNets,其中第一个块的步幅为2。考虑到前几个块带来的大计算量,我们将Stand-Alone Axial-ResNets中所有的块中轴向跨距m设为15。注意力头数全部都是N=8[68]。为了不必小心翼翼地初始化、、、、、,我们在所有的注意力层都使用了批量归一化[43]。

表1总结了ImageNet的结果。基线ResNet-50[32](由[68]实现)和[68]的Conv Stem+Attention模型也列在表中。在conv-stem设置中,在[68]的注意力层中增加批量归一化可稍稍提高0.3%的性能。我们提出的位置敏感自我注意(Conv-Stem+PS-attention)只带来极少的额外计算量,进一步提高了0.4%的性能。我们提出的Conv-Stem + Axial-Attention和[68]的Conv Stem+Attention在准确度上相当,而参数更少、计算量更少。在全注意力模型的比较中,我们的全注意力模型比[68]的全注意力模型高0.5%,参数效率高1.44倍,计算效率高1.09倍。

和[68]一样,我们设置不同的网络宽度,即Axial-ResNets {S,M,L,XL},以探索准确度、参数量和计算量(也就是M-Adds)的平衡。如图3所示,在所有设置下,我们提出的Conv-Stem + PS-Attention和Conv-Stem + Axial-Attention已经好于ResNet-50[32,68]和[68]提出的注意力模型(包括Conv-Stem和全注意力模型)。我们提出的Full-Axial-Attention模型进一步取得最优的准确度-参数量和准确度-计算量的平衡。

图3 在ImageNet分类任务上,对比不同准确率下的参数量和乘-加计算量。我们的位置感知自注意力(ConvStem+PS-Attention)和轴向注意力(ConvStem+Axial-Attention),在各种网络宽度(也即通道数)下,是一直好于ResNet50[32,68]和注意力模型[68](包括卷积柄+注意力,以及全注意力网络)。在参数量和乘加计算量上,我们的全-轴向注意力模型都是最佳的

4.2 COCO

将ImageNet预训练的Axial-ResNet模型的各种变体(通道数不同的各种版本)转化为相应的Axial-DeepLab,然后用于全景分割任务。

我们首先在Challengeng COCO数据集[59]上证明了我们的Axial-DeepLab的有效性,该数据集包含各种尺度的物体(从小于32×32的到大于96×96的)。

验证集:    在表2中,我们报告了在验证集上的结果,因为我们的方法属于自底向上全景分割方法的一类,因此我们和其它的自底向上全景分割方法进行比较。如表所示,我们的单尺度Axial-DeepLab-S比DeepLab[93]高出8%PQ,比多尺度SSAP[29]高出5.3%PQ,比单尺度Panoptic-DeepLab高出2.1%PQ。有趣的是,我们的单尺度Axial-DeepLab-S比多尺度的Panoptic-DeepLab高出0.6%PQ,而且参数效率高3.8倍,计算(M-Adds)效率高27倍。如果增加主干网络的容量的话(也就是加大通道数),性能还可继续不断地提高。特别的,我们的多尺度Axial-DeepLab-L取得了43.9%的PQ,比Panoptic-DeepLab[20]高出2.7%PQ。

Test-dev测试集:    如表3所示,当使用更大的主干时,Axial-DeepLab的PQ就会越高。我们的多尺度Axial-DeepLab-L达到44.2%PQ,比DeeperLab[93]高9.9%PQ,比SSAP[29]高7.3%PQ,比Panoptic DeepLab[20]高出2.8%PQ,在自底向上方法中达到最先进水平。表中还列出了最先进的自顶向下方法,以供参考。

尺度鲁棒性测试:    为了验证我们的方法能学习远程关系,我们除了标准测试还实验了尺度鲁棒性测试。我们按照标准设置训练了Panoptic-DeepLab(X-71)和我们的Axial-DeepLab-L,不过我们将测试集的分辨率设置的和训练集不同(out-of-distribution resolutions),具体是将图像缩放到不同分辨率。图4是Panoptic DeepLab分别在PQ、PQ(thing)和PQ(stuff)的相对提升。在大分辨率图像上测试时,Axial-DeepLab的性能有大的提升(30%),这证明了它比卷积能更好的编码远程关系。此外,在小图像上,Axial DeepLab性能提高了40%,这表明轴向注意力对尺度变化更具鲁棒性。

图4 在COCO验证集上的尺度鲁棒测试。在极端分辨率下,Axial-DeepLab的性能优势就越大。图中的x坐标轴,比率4.0表示分辨率为4097x4097

4.3 Mapillary Vistas

我们在大规模的Mapillary Vistas上评估我们的Axial-DeepLab。我们只报告了验证集的结果,因为测试集服务器不可用。

验证集:    如表4所示,我们的Axial-DeepLab-L在单尺度和多尺度情况下都优于所有最先进的方法。我们的单尺度Axial-DeepLab-L的性能比之前最好的单尺度Panoptic-DeepLab(X-71)[20]好2.4%。多尺度模型的对比中,我们的轻量级Axial-DeepLab-L比Panoptic DeepLab(Auto DeepLab XL++)性能更好,不仅在全景分割(0.8%PQ)和实例分割(0.3%AP)上,而且在语义分割(0.8%mIoU)上,而Auto-DeepLab[60]的网络结构是在语义分割上搜索来的。另外,据我们所知,我们的Axial-DeepLab-L取得了最好的单模型语义分割结果。

4.4 Cityscapes

验证集:    表5(a)报告了在Cityscapes验证集上的结果。没有使用额外的数据的时候,也就是说,只使用Cityscapes的精标注数据,我们的Axial-DeepLab取得了65.1%的PQ,比当前的最先进的自底向上方法Panoptic-DeepLab[20]高1%,比基于proposal的AdaptIS[80]高3.1%。当使用额外数据(即Mapillary Vistas[65])训练时,我们的多尺度Axial-DeepLab-XL取得了68.5%PQ,比Panoptic-DeepLab[20]高1.5%,比Seamless[71]高3.5%。我们的实例分割和语义分割结果分别比Panoptic DeepLab[20]好1.7%和1.5%。

测试集:    表5(b)是在测试集上的结果。在没有额外数据的情况下,Axial DeepLab XL达到62.8%的PQ,取得了最先进结果。在Mapillary Vistas预训练的情况下,进一步实现了66.6%的PQ,39.6%的AP,84.1%的mIoU。请注意,Panoptic DeepLab[20]在测试集进行推理时采用了输出步幅8的技巧,这样他们的M-Adds和我们的XL模型具有可比性。

4.5 消冗研究

我们在Cityscapes验证集上做消冗研究。

位置敏感和轴向注意力的重要性:    在表1里我们是在ImageNet上对注意力模型实验。而在本节中,我们将它们用到Cityscapes分割任务上。如表6所示,所有的变体都要比ResNet-50[32]好。位置敏感的注意力模型要比先前的最先进的自注意力模型[68]要好,这与表1中在ImageNet上的结果是一致的。然而,所不同的是,在表1中,在ImageNet上,使用轴向注意力与使用位置敏感的自注意力性能是相当的(笔者注:应该是将本文方法和[68]方法比较,表中axial-attention指本文方法,attention指[68]方法),而在表6中,轴向注意力对3个分割任务(语义、实例、全景)都带来超过1%的提升,没有使用ASPP,参数更少,计算量(M-Adds)更少,这表明轴向注意力建模远程上下文关系的能力显著提高在大输入图像分割任务的性能。

轴向注意力跨距的重要性:    在表7中,我们在不使用ASPP的情况下改变跨距m(也就是轴向块中局部区域的空间范围)。我们观察到,使用更大的跨距,性能就越高,而只带来极少的计算量增量。

5. 结论和讨论

本文证明了位置敏感的轴向注意力在图像分类和分割任务中是有效的。在ImageNet上,我们通过堆叠轴向注意力块构建的Axial-ResNet在全注意力模型中取得最先进水平。我们将Axial-ResNet进一步转换为Axial-DeepLab用于自底向上分割任务,在几个基准数据集(COCO、Mapillary Vistas和Cityscapes)上也取得最先进水平。我们希望这些好的结果能够证明,轴向注意力是现代计算机视觉模型的有效组成部分。

我们的方法与解耦卷积(decoupled convolution)[44]相似,其是将深度卷积[78,36,21]分解为列卷积和行卷积。这种操作理论上也可获得大的感受野,但是卷积本身的模板匹配特性限制了多尺度关系建模的能力。我们方法还与可变形卷积相似[24, 100,28],其是每个点动态地关注图像上的几个点。然而,可变形卷积没有利用关键依赖的位置偏置,也没有利用基于上下文的关系。此外,轴向注意力可以密集地、高效地沿着高度方向和宽度方向传播信息。

虽然我们的轴向注意力模型节省了M-Adds,但它比卷积对应的模型慢,正如[68]所观察到的。这是由于目前各种加速器上缺乏专门的内核。如果社区认为轴向注意力是一个合理的方向,这一点可能会得到改善。

附录A  运行时间

在本节中,我们用一个常用设置来介绍Conv-Stem Axial-ResNet-L:输入224x224的图像,批量大小设为1,在V100GPU上运行5次。耗时包括数据标准化,以及最后映射到1000logits上。我们的模型耗时16.54ms。我们还对比了我们手头上的一些Tensorflow实现的常用模型,它们的flops与我们的模型是相当的。为了提供更多的背景信息,我们引用了[7]中的一些实验结果,[7]用的是Titan X Pascal,但是Pytorch代码是经过更多优化的。我们的耗时与ResNeXt-101(32x4d)、SE-ResNet-101、ResNet-152和DenseNet-201(k=32)是大致相同的。

请注意,我们是直接用为TPU计算而优化过的代码来测试。从经验上,生成的图包含从NCHW到NHWC的转换。这一效果不利于Xception-71,因为它包含深度可分离卷积。这一点进一步的优化可导致更快的速度。

我们观察到,ConvStem Axial-ResNet-L比ConvStem Stand-Alone-L[68]运行得更快,尽管我们将一层分成两层。这是因为我们的轴向注意力更好地利用了现有的核:

—宽度方向的注意力在高度方向上是可以并行的,也就是说,这是个大批量的1维的行操作,批量大小就是输入的高度。

—轴向注意力避免了提取2维的内存块:包含填充、分拆和拼接,而这些是不利于加速的。

附录B 轴向解码器

Axial-DeepLab使用的是Panoptic-DeepLab[20]的双卷积解码器。本节我们探讨一种使用单个轴向注意力的解码器的设计。在轴向解码器中,我们在每个上采样阶段使用一个轴向注意力块。图5所示是输出步长8到输出步长4的轴向解码器。我们使用了3个这样的块,类似于Panoptic-DeepLab[20]中的3个5x5卷积。

图5 轴向解码器块。我们在轴向注意力块上增加上采样和编码器特征

输出步幅和轴向解码器的重要性:表9是不同输出步幅以及是否用轴向解码器替换双卷积解码器的实验。如表所示,我们的模型对输出步幅是鲁棒的,使用轴向解码器还是一样对输出步幅鲁棒。我们设计的轴向解码器和双卷积解码器效果一样好。

附录C 对COCO数据的可视化 

图6是对COCO验证集上一些数据的全景分割结果的可视化。与Panoptic-DeepLab(Xception-71)[20]相比,我们的Axial-DeepLab-L显示了对遮挡的鲁棒性。

图6 COCO验证集一些结果的可视化。 Axial-DeepLab显示了对遮挡的鲁棒性。 在第1行和第4行中,Axial-DeepLab准确地捕捉到了被遮挡的左腿和遥控器的线,这些甚至连真值都没能捕捉到。 最后一行中,Axial-DeepLab准确的区分了两个人,而真值将他们视为一个人。

在图7和图8中,我们在COCO验证集上可视化了我们的Axial-DeepLab-L的注意力映射。我们可视化了一个低级块(stage3的block2)和一个高级块(stage4的block3),它们分别是在输出步幅16的设置下,分辨率为65x65的第一个块和最后一个块。我们发现,在多头注意力设置中,一些注意力头学习到关注局部细节,而另一些则关注远程上下文。我们还发现,一些注意力头能捕捉位置信息,而另一些注意力头学会关联语义概念。

图7 stage3的block2的注意力映射。 我们选取一行像素,并可视化8个注意力头对它们的列注意力(高度方向)。 然后,我们选取一列像素,并可视化它们的行注意力。 蓝色像素是我们选取的queries,红色像素表示相应的注意力权重。 我们注意到,column head1会关联到人的头部上,而column head4只关联草地区域。 row head 6 关注相对局部区域,而row head 5 关注整张图像区域。
图8 stage4的block3。相比图7中的结果,它们会更关注远程的上下文,尽管两者都具有全局感受野。

在图9中,我们根据Panoptic DeepLab[20]中定义的三个训练损失函数,比较了Axial DeepLab和Panoptic DeepLab[20]。我们观察到轴向DeepLab能够更好地拟合数据,尤其是在预测物体中心偏移量的任务上。这也证明了我们的位置感知注意力设计的有效性,以及轴向注意力的远程建模能力。

图9 COCO数据上的训练损失。 相比Panoptic-DeepLab[20],我们的带有位置感知的轴向注意力的Axial-DeepLab能更好地拟合数据, 尤其是在物体中心偏移量预测任务上,该任务需要准确而远程的位置信息。

附录D原始数据

在正文图3中,我们展示了ImageNet分类任务上不同准确率下参数量和乘加计算量的对比,在这里我们做个补充,我们在表10中展示我们模型的性能。

在正文的图4中,我们展示了我们的Axial-DeepLab-L相比Panoptic DeepLab(Xception-71)在COCO数据集的尺度鲁棒测试下的提升。这里,我们在图10中展示2个模型的性能。

图10 COCO验证集的尺度鲁棒测试

你可能感兴趣的:(横着注意力+竖着注意力)