深度学习之图像分类(十四)CAT: Cross Attention in Vision Transformer详解
IPSA和CPSA的处理流程、维度变换细节 |
FLOPs的计算方法、以及flops和划分的patch数目以及patch的维度计算关系 |
IPSA如何进行local attention、CPSA如何进行globe attention |
CAT的代码详细注释---需要学习完Transformer |
TNT、swin transformer、crossViT |
CAT: Cross Attention in Vision Transformer
Abstract
由于Transformer在自然语言处理(NLP)中得到了广泛应用,人们已经意识到Transformer在计算机视觉(CV)中的潜力,并且已经激发了许多新的方法。然而,在图像进行标记后,用图像块替换单词标记所需的计算量(例如,ViT)非常庞大,这成为模型训练和推理的瓶颈。在本文中,我们提出了一种新的Transformer中的注意机制,称为交叉注意力(Cross Attention),该机制在图像块内部而不是整个图像中进行交替注意以捕捉局部信息,并在从单通道特征图中分割的图像块之间应用注意力以捕捉全局信息。这两种操作的计算量都比Transformer中的标准自注意力小。通过交替在块内和块间应用注意力,我们实现了交叉注意力,以在较低的计算成本下保持性能,并构建了一种称为Cross Attention Transformer(CAT)的分层网络,用于其他视觉任务。我们的基础模型在ImageNet-1K上取得了最新水平,并在COCO和ADE20K上改进了其他方法的性能,说明我们的网络具有作为通用骨干结构的潜力。
1 Introduction
随着深度学习的发展和卷积神经网络(CNN)的应用[1],计算机视觉任务取得了巨大的进步。自2012年以来,CNN在很长一段时间内一直主导着计算机视觉,它作为各种视觉任务中的关键特征提取器,以及其他任务中的任务分支编码器。各种基于CNN的网络[2–11]具有不同的改进和应用,各种下游任务也有这些多种方法,比如目标检测[12–20]、语义分割[21–28]等。
最近,作为一种新的网络结构,Transformer[29]在自然语言处理(NLP)中取得了显著的成果。由于其出色的提取全局信息的能力,它还解决了诸如循环神经网络(RNN)[30]和长短时记忆网络(LSTM)[31]等(RNN和LSTM---难以并行化的序列模型)难以并行化的序列模型的问题,使得自然语言处理领域取得了重要的飞跃,同时也激发了计算机视觉任务的发展。
最近的研究[32–41]将Transformer引入计算机视觉作为图像提取器。然而,在自然语言处理中,文本序列的长度是固定的,这导致了Transformer处理图像的能力下降,因为不同任务的输入分辨率是可变的。在使用Transformer处理图像时,(IGPT的做法是将像素类比为NLP中的单词,这将导致计算量巨大的增加。)一种朴素的方法是将每个像素视为一个类似于单词标记的全局注意力标记。iGPT[42]证明了这样做带来的计算量是巨大的。一些工作(例如ViT,iGPT)将一个区域内的一组像素视为一个标记(将16x16的patch看做一个token),这在一定程度上减少了计算量。然而,随着输入大小的增加(公式1),计算复杂性急剧增加,而这些方法生成的特征图具有相同的形状(图1(b)),使得这些方法不适合用作后续任务的骨干结构。(特征图大小固定,可扩展性较低)
图1:分层网络。 (a) 基于CNN的分层网络,不同阶段生成具有不同尺度的特征。 (b) 基于Transformer(例如,ViT)的分层网络,所有特征在形状上都相同。 (c) CAT(我们的)的分层网络,具有CNN层次结构的特征。
(在patch内部先进行注意力,然后将获取到patch内全部像素相关信息的token进行拼接,然后再在拼接的代表每个patch的token之间进行attention,此时获取的全局token包含全局的信息,然后将包含全局信息的token与包含每个patch局部信息的token之间进行attention,这种attention的方式称为 cross attention)
在本文中,(CNN:局部特征提取,将局部特征加权计算后拼成新的特征图,由局部特征--->全局特征)我们受到了CNN局部特征提取能力的启发,采用了在一个图像块内部的像素之间进行注意力(在patch内部先进行注意力,然后将获取到patch内全部像素相关信息的token进行拼接,然后再拼接的代表每个patch的token之间进行attention,此时获取的token包含全局的信息,然后将包含全部信息的token与包含每个patch局部信息的token之间进行attention,这种attention的方式称为 cross attention),以模拟CNN的局部特征提取方式,将随着输入大小呈指数增长的计算减少到与图像块大小呈指数关系的程度。(当前的计算复杂度只与patch块的大小相关)同时,如图3(b)所示,为了考虑图片的整体信息提取和传递,我们设计了一种在单通道特征图上执行注意力的方法。与对所有通道进行注意力相比,计算量大幅减少,如公式1和图3所示。通过交替进行图像块内部的内部注意力和单通道特征图的注意力,实现了交叉注意力。(多通道的局部token与单通道的全局token进行cross attention)我们可以利用交叉注意力构建一个强大的骨干结构,生成不同尺度的特征图,满足下游任务对不同颗粒度特征的需求,如图1所示。我们引入全局注意力,而不增加计算或仅略微增加计算,这是一种更合理的方法,可以联合Transformer和CNN的特征。
我们的基础模型在ImageNet-1K上实现了82.8%的top-1准确率,与当前的基于CNN和Transformer的先进网络相媲美。同时,在其他视觉任务中,我们的CAT作为目标检测和语义分割方法的骨干结构可以提高它们的性能。
Transformer和CNN的特性互补,将它们更有效地、完美地结合是我们的长期目标,而我们提出的CAT是朝着这个方向迈出的一步。希望在这个方向上会有更好的发展。
2 Related work
CNN(卷积神经网络)/基于CNN的网络具有(归纳偏执)权重共享、平移、旋转不变性和局部性等特点,在过去的十年中,在计算机视觉领域取得了巨大的成就,取代了多层感知器(MLP),成为视觉任务中的标准网络。作为在计算机视觉中取得巨大成功的第一个CNN网络,AlexNet为后来基于CNN的网络的发展奠定了基础,[3, 8–11, 43]用于性能改进已成为视觉任务中的骨干选择。为提高效率而设计的Inceptions[4, 5, 44–46]、MobileNets[6, 47, 48]和ShuffleNets[7, 49]在需要推理速度的任务中也是替代选择。
在基于Transformer的网络中的全局注意力。Transformer最初在自然语言处理中用于机器翻译,在那里核心的多头自注意力(MSA[29])机制对于在多个层次上提取单词之间关系的特征至关重要。作为最早的几个基于Transformer的骨干结构,ViT[33]和Deit[38]将图像划分为图像块(图像块大小为16×16)。一个图像块被扁平化为一个标记,并引入了CLS-Token[50]进行分类。CvT[35]和CeiT[51]都引入了卷积层来替代QKV[29]的线性投影。CrossViT[34]通过将图像划分为不同大小的图像块,将不同粒度的全局特征集成到两个分支中。然而,这些方法将所有图像块放在一起进行MSA,只关注不同图像块之间的关系,随着输入大小的增加,计算复杂性呈指数增加,如公式1所示,这很难应用于需要大分辨率输入的视觉任务。
在基于Transformer的网络中的局部注意力。图像块内部信息的关系对于视觉任务[52, 53]至关重要。最近,TNT[37]将每个图像块划分为更小的图像块。通过提出的TNT块,捕捉了全局信息和图像块内部信息。Swin[36]将每个图像块视为窗口,以提取图像块的内部相关性,并使用平移窗口来捕捉更多特征。然而,这两种方法都有问题。首先,在TNT[37]中,为了结合全局信息交互和局部信息交互,计算量的增加不能被低估。其次,在[36]Swin中,局部信息交互与相邻图像块之间的交互缺乏全局信息交互。我们提出了一个交叉图像块自注意力块,可以有效地保持全局信息交互,同时避免随着输入分辨率的增加而导致的计算量的巨大增加。
层次网络和下游任务。Transformer已成功用于视觉任务[54–57]和NLP任务[50, 58–63]。然而,由于典型Transformer中输入和输出的一致形状,要实现与CNN-based网络[8, 9, 19, 48]类似的层次结构在下游任务中是困难的。FPNs[14, 20, 64, 65]与ResNet[8]结合已成为目标检测中的标准范例。在语义分割[21, 23, 28]中,金字塔形的特征被用来提高性能。最近的PVT[32, 36]和Swin[36]降低了不同阶段的特征分辨率,类似于ResNet[8],这也是我们采用的方法。
3 Method
3.1 Overall architecture
我们的方法旨在将图像块内部的注意力和图像块之间的注意力相结合,并通过堆叠基本块构建一个层次网络,这可以简单地应用于其他视觉任务。如图2(a)所示,首先,我们将输入图像缩小到H1 = H/P,W1 = W/P(其中P = 4在我们的实验中)(之后经过liner projection),并通过参考ViT[33]中的图像块处理模式,通过图像块嵌入层(patch embedding)将通道数增加到C1。然后,使用几个CAT层在不同尺度上进行特征提取。
图2: (a) CAT架构,第三阶段的CAB数量随模型大小而变化。 (b) 交叉注意力块(CAB),堆叠IPSA和CPSA,都带有LN[66],MLP和快捷连接[8]。
经过上述预处理,输入图像进入第一阶段。此时,图像块的数量为H1/N × W1/N,图像块的形状为N × N × C1(其中N是图像块嵌入层后的图像块大小)。第一阶段输出的特征图的形状为H1 × W1 × C1,表示为F1。然后,进入第二阶段,图像块投影层执行空间深度操作,该操作使形状为2 × 2 × C的像素块从形状为2 × 2 × C变为形状为1 × 1 × 4C,然后通过线性投影层投影到形状为1 × 1 × 2C。在接下来的阶段通过几个交叉注意力块,生成形状为H1/2 × W1/2 × C2的F2,特征图的长度和宽度可以减小一倍,维度增加一倍,类似于ResNet[8]中的操作,也是Swin[36]中的实践。经过四个阶段后,我们可以得到{F1,F2,F3,F4},四个不同尺度和维度的特征图。与典型的基于CNN的网络[8, 9]一样,可以为其他下游视觉任务提供不同粒度的特征图。
3.1.1 Inner-Patch Self-Attention Block
在计算机视觉中,每个像素都需要一个特定的通道来表示其不同的语义特征。类似于NLP中的单词标记,理想情况是将特征图的每个像素都视为一个标记(例如,ViT,DeiT),但计算成本太高。正如公式1所示,计算复杂性随着输入图像的分辨率呈指数增长。例如,在传统的RCNN系列[12, 13, 67, 68]方法中,输入的短边至少为800像素,而YOLO系列[16, 17, 69]的方法也需要超过500像素的图像。大多数语义分割方法[21, 22, 25]也需要边长为512像素的图像。计算成本至少比预训练阶段的224像素高5倍。
FLOPsMSA = 4HWC2 + 2H2W2C (1)
受到CNN局部特征提取特性的启发,我们将CNN中的卷积方法的局部性引入Transformer,以在每个图像块中进行像素级自注意力,称为图像块内自注意力(Inner-Patch Self-Attention,IPSA),如图3(a)所示。我们将一个图像块视为一个注意力范围,而不是整个图片。同时,Transformer可以根据输入生成不同的注意力图,这与具有固定参数的CNN相比具有显著优势,类似于卷积方法中的动态参数,这在[用于实例分割的条件卷积70]中被证明是有益的。[37]揭示了像素之间的注意力同样至关重要。我们的方法显著减少了计算量,同时考虑了图像块内像素之间的关系。计算公式如下:
FLOPsIPSA = 4HWC2 + 2N2HWC (2)
其中N是IPSA中的图像块大小。与标准Transformer中的MSA相比,计算复杂性从与H × W的二次相关性(公式1)降低到与H × W的线性相关性。假设H,W = 56,C = 96,N = 7,根据公式1,FLOPsMSA ≈ 2.0 G,而根据公式2,FLOPsIPSA ≈ 0.15 G,大大减少了计算量。
图3: IPSA和CPSA的流程。 (a) IPSA: 展开所有通道的输入为2×2,然后堆叠它们,在IPSA块之后,重新调整到原始形状。 (b) CPSA: 将单通道输入展开为2×2的图像块,然后堆叠它们,在CPSA块之后,重新调整到原始形状。
3.1.2 Cross-Patch Self-Attention Block
将注意力机制引入像素之间只能确保捕捉一个图像块内的像素之间的相互关系,但整个图片的信息交换也非常重要。在基于CNN的网络中,通常通过堆叠卷积核来扩大感受野。提出了扩张/空洞卷积[24](深度可分离卷积和空间可分离卷积)以获得更大的感受野,实际上期望最终的感受野扩展到整个图片。Transformer天然具备捕捉全局信息的能力,但像ViT[33]和Deit[38]这样的尝试最终并不是分辨率最佳。
每个单通道特征图天然具有全局空间信息。我们提出了Cross-Patch Self-Attention,将每个通道的特征图分开,并将每个通道划分为H/N × W/N个图像块,使用自注意力获取整个特征图的全局信息。这类似于Xception[46]和MobileNet[6]中使用的深度可分离卷积。我们方法的计算可以如下计算:
FLOPsCPSA = 4N2HWC+ 2(HW/N)2C (3)
其中N是CPSA中的图像块大小,H、W分别表示特征图的高度和宽度。计算成本比ViT(公式1)和其他基于全局注意力的方法要少。同时,如图2(b)所示,我们结合了MobileNet[6]的设计,堆叠IPSA块和CPSA块,以提取和整合一个特征图中像素之间和一个特征图中图像块之间的特征。与Swin[36]中手动设计的平移窗口相比,Swin难以实现,并且很难捕捉全局信息,而我们的方法更为合理且易于理解。按照上一节的假设,FLOPsCPSA约为0.1 G,远少于MSA的2.0 G。
多头自注意力机制在[29]中被提出。在NLP中,每个头可以注意到不同的语义信息,例如单词之间的关系。在计算机视觉中,每个头可以注意到图像块之间不同的语义信息,这类似于基于CNN的网络中的通道(多头机制类似于CNN中的通道)。在CPSA中,我们将头的数量设置为图像块的大小,使得一个头的维度等于图像块的大小,但这对性能没有帮助,如表5所示。因此,在我们的实验中,默认设置为单头。
表5: 在ImageNet-1K上使用CAT-S架构进行的CPSA多头、CAB中第二个IPSA块中的平移窗口、以及图像块嵌入层中的切片或卷积方法的消融研究。
(多通道进行注意力机制时使用相对位置编码,但是在单通道使用了绝对位置编码)
位置编码。我们采用了相对位置编码,参考了[36, 71, 72]中的IPSA,而对于在整个单通道特征图上执行自注意力的CPSA,我们对嵌入到图像块嵌入层中的特征添加了绝对位置编码,可以表示为:
y = Patch.Emb(xinput) (4)
ytemp = IPSA(y + ab.pos.) [] (5)
youtput = IPSA(CPSA(ytemp)) [] (6)
其中ab.pos.表示绝对位置编码,Patch.Emb表示表1中的图像块嵌入层。绝对位置编码在CPSA中对于提高性能是有用的,表6中报告了相关结果。
表6: 在CAT-S架构上,关于CPSA中的绝对位置编码和自注意力中的dropout的消融研究,使用FCOS[15]在COCO 2017上进行了1x计划,使用ADE20K上80k迭代的Semanticc FPN[90]。attn.d:自注意力的dropout。abs.pos.:绝对位置编码。
3.1.3 Cross Attention based Transformer
交叉注意力块由两个图像块内自注意力块和一个图像块间自注意力块组成,如图2(b)所示。CAT层由多个CAB组成,网络的每个阶段由不同数量的层和一个图像块嵌入层组成,如图2(a)所示,CAB的流程如下:
ˆytemp1 = IPSA(LN (ˆyn−1)) + ˆyn−1 (7)
ˆytemp2 = MLP(LN (ˆytemp1)) + ˆytemp1 (8)
ˆytemp3 = CPSA(LN (ˆytemp2)) + ˆytemp2 (9)
ˆytemp4 = MLP(LN (ˆytemp3)) + ˆytemp3 (10)
ˆytemp5 = IPSA(LN (ˆytemp4)) + ˆytemp4 (11)
ˆyn = MLP(LN(ˆytemp5)) + ˆytemp5 (12)
其中,ˆytempi是带有LN的一个块(例如IPSA、MLP)的输出。我们比较了[33]中图像块嵌入层的卷积,其中卷积核大小设置为P,步长也为P,以及将输入切片的方法[16],结果见表5,两者性能相同。我们的默认设A置是前者。根据阶段3中CAB的数量和图像块投影层的维度,设计了三个不同计算复杂性的模型,分别是CAT-T、CAT-S和CAT-B,其计算量分别为1×、2×和3×。表1详细说明了配置。
表CATs的详细配置。Down.rate表示每个阶段的下采样速率。R表示特定层的dowm采样率。
表5: 在ImageNet-1K上使用CAT-S架构进行的CPSA多头、CAB中第二个IPSA块中的平移窗口、以及图像块嵌入层中的切片或卷积方法的消融研究。
4 Experiment
我们分别在ImageNet-1K[73]、COCO 2017[74]和ADE20K[75]上进行图像分类、目标检测和语义分割实验。接下来,我们将CAT架构与现有技术架构在这三个任务上进行比较,然后报告一些我们在CAT中采用的设计的消融实验。
4.1 Image Classification
细节:对于图像分类,我们在ImageNet1K[73]上使用单一裁剪报告top-1准确率,该数据集包含来自1000个类别的128万训练图像和5万验证图像。我们的实验设置主要遵循[38]。我们采用批量大小为1024,初始学习率为0.001,权重衰减为0.05。我们使用AdamW[76]优化器、余弦衰减学习率调度器和20个周期的线性预热进行300个周期的模型训练。在训练中使用了随机深度[77],分别对三种变体架构采用了0.1、0.2和0.3的比率,而在CAB的自注意力中采用了0.2的dropout[78]比率以防止过拟合。我们使用了[38]中的大多数正则化策略和数据增强,这使得我们的结果更具可比性和说服力,与[36]类似。
结果表2中呈现了我们的实验结果,表明我们的CAT-T在FLOPs减少65%的情况下,能够达到80.3%的top-1精度,比ResNet101[8]高。同时,CAT-S和CAT-B在分辨率为224×224的图像上的top-1分别为81.8%和82.8%。这样的结果与表中现有技术的结果相媲美。例如,与Swin-T[36]相比,其计算相似,我们的CAT-S提高了0.5%。特别是,我们的方法比Swin[36]中的平移操作更具捕捉图像块之间关系的能力。Swin-T(w. shifted)提高了1.1%的top-1精度,而CAT-S超过了1.6%。
4.2 Object detection
细节:对于目标检测,我们在COCO 2017[74]上进行实验,使用mAP作为度量标准,该数据集包含来自80个类别的118k训练图像、5k验证图像和20k测试图像。我们在一些框架上进行实验以评估我们的架构。实验中使用批量大小为16,初始学习率为1e-4,权重衰减为0.05。我们采用了AdamW[76]优化器、1x计划和NMS[80]。其他设置与MMDetection[81]相同。注意采用随机深度[77]率为0.2以防止过拟合。关于多尺度策略,我们使用从480到800的间隔为32的随机选择的短边尺度进行训练,而长边尺度小于1333,与[54, 82]相同。
结果:如表3所示,我们在一些基于锚点和无锚点的框架中使用了CAT-S和CAT-B作为骨干网络,两者性能更好且计算成本相当或更低。CAT-S通过多尺度策略提高了FCOS[15]达3.4%,RetinaNet[13]为3.7%,Cascade R-CNN[85]为4.8%。而对于实例分割,我们使用了MASK R-CNN[83]框架,其中CAT-S的掩码mAP提高了4.2%。我们实验的所有方法都比原始方法性能更好,证明了我们的CAT在特征提取方面具有更强的能力。
表CAT与其他主干在COCO检测上的各种方法的比较。−表示采用多尺度策略进行训练。FLOPs是在800 × 1280上评估的。
4.3 Semantic Segmentation
细节:对于语义分割,我们在ADE20K[75]上进行实验,该数据集包含20k张训练图像、2k张验证图像和3k张测试图像。设置如下,初始学习率为6e-5,总共进行160k和80k次迭代,批量大小为16,权重衰减为0.05,预热迭代次数为1500。我们在Semantic FPN[86]框架下进行实验,输入为512×512,并使用了MMSegmentation[87]的基本设置。注意,在CAT的训练过程中使用了随机深度[77]率为0.2。
结果:如表4所示,我们采用CAT-S和CAT-B作为骨干网络,结合Semantic FPN[90]框架。Semantic FPN在CAT-S和CAT-B的支持下取得了更好的性能,特别是在160k迭代和CAT-B的情况下,我们取得了44.9%的mIoU,相较于将ResNet101[8]作为骨干网络的情况提高了4.2%,使Semantic FPN的性能达到了与其他方法相当的水平,而在80k迭代中,结果提高了4.8%,说明我们的架构比ResNet[8]更适合作为骨干网络。
表4.ADE20K上的语义分割性能。指示模型已在ImageNet-22k上进行预训练。表明经过80k次迭代训练。FLOPs是在1024 × 1024上计算的
4.4 Ablation Study
在本节中,我们报告了在设计架构和在ImageNet-1K[73]、COCO 2017[74]和ADE20K[75]上进行实验时所做设计的一些消融实验的结果。
图块嵌入函数:我们比较了图块嵌入层中的嵌入函数,包括卷积方法和[16]中的方法。卷积方法使用卷积层,卷积核大小为4×4,步长为4,将输入的分辨率降低到原始图像的1/4;后者从H×W×H×C切片输入为H/S×W/S×SC,其中S在我们的模型中设为4,以实现与前者相同的效果。表5中的结果表明这两种方法的性能相同。为了更好地与其他工作[36]进行比较,我们选择卷积方法作为默认设置。
多头和偏移窗口:多头机制是在[29]中提出的,多头表示不同语义特征之间的关系。我们将每个CPSA中的头数设置为与图块大小相等(CPSA的数量等于patch_number的数量),这对性能无益,如表5所示。为了研究Swin[36]中的偏移窗口,我们还在CAB的第三个块中进行了实验,结果显示偏移操作在我们的架构中表现并不更好。
CPSA的绝对位置和自注意力中的dropout:我们对CPSA的绝对位置编码进行了消融研究,结果表明它在三个基准上提高了性能。为了更好地训练,我们在CPSA的自注意力中采用了dropout[78],并设置了0.0和0.2的率。0.2的率达到了最佳性能,说明CPSA中存在一些过拟合。所有结果见表6。
表5:在ImageNet-1K上使用CAT-S架构进行的关于CPSA中的多头、CAB中第二个IPSA块的偏移窗口以及图块嵌入层中切片或卷积方法的消融研究。
表6:在使用CAT-S架构的三个基准上对CPSA的绝对位置编码(abs.pos)和自注意力中的dropout(attn.d)进行的消融研究。在COCO 2017上使用1x计划的FCOS[15]和在ADE20K上进行80k次迭代的Semantic FPN[90]。attn.d:自注意力的dropout。abs.pos.:绝对位置编码。
5 Conclusion
在本文中,提出了交叉注意力(Cross Attention)以更好地结合CNN中的局部特征提取和Transformer中的全局信息提取,构建了一个强大的骨干网络,即CAT。CAT可以生成类似于大多数基于CNN的网络的不同尺度的特征,并且还可以适应其他视觉任务的不同输入尺寸。CAT在各种视觉任务数据集上(例如ImageNet-1K[73]、COCO 2017[74]、ADE20K[75])取得了领先水平的性能。关键在于我们在特征图块内部和单通道特征图上交替注意力,而几乎不增加计算量以捕捉局部和全局信息。我们希望我们的工作将是将CNN和Transformer集成为多领域方法的方向中的一步。