Swin Transformer: 使用滑动窗口的分层视觉transformer

论文地址:https://arxiv.org/abs/2103.14030

官方GitHub:https://github.com/microsoft/Swin-Transformer

目录

0、摘要:

1、引言

2、相关工作

3、方法

3.1、总体架构

3.2、基于 Shifted Window的self-attention

3.3、架构变体

4、实验结果

5、结论


0、摘要:

本文提出了一个新的视觉transformer,命名为Swin Transformer,可以作为视觉任务的通用backbone。将transformer从NLP迁移到CV主要障碍来自两个领域之间的差异性,如:视觉实体的尺度变化多端、图像中的像素相比于文本中的单词具有更高的分辨率。为了解决这些差异,我们提出了一个分层transformer,其表示是使用滑动窗口计算而得的。滑窗机制将self-attention的计算限制在一个局部窗口中,同时允许跨窗口连接,从而带来了很高的效率。而分层结构可以灵活建模不同的尺度,且计算复杂度与图像尺寸线性相关。Swin Transformer的这些特性,使得它能够处理各种cv任务,如图像分类(ImageNet-1K上达到86.4的top-1准确率)、目标检测(COCO test-dev上达到58.7 box AP 和51.1 mask AP )、语义分割(ADE20K val上达到53.5 mIoU),均超过了先前的模型,达到新的SOTA,这也表明了基于transformer的模型具备了作为视觉backbone的潜力。

1、引言

CV建模一直被CNN主导,从AlexNet及其在ImageNet分类挑战赛上革命性的性能开始,CNN架构通过更大的尺度、更多的连接、更复杂的卷积形式,变得越来越强大。随着CNN作为各种视觉任务的backbone,其架构不断改进,其性能不断提升,随之而来的是整个领域的进步。

与此同时,在NLP领域网络结构的进化采用了一种不同的方式,其当前的主流架构被transformer所主导。专为序列建模和转换任务而设计的Transformer以其对数据中的远程依赖关系建模而著称。其在NLP领域的巨大成功,使得研究人员开始研究它对计算机视觉的适应,最近在一些任务上展示了很有前途的结果,特别是图像分类和联合视觉语言建模。

在本文中,我们尝试扩展Transformer的能力,使其能够作为视觉任务的通用backbone,就像它在NLP中、CNN在CV中所扮演的角色那样。我们观察到,将transformer在NLP中的高性能迁移到视觉领域时一个巨大的挑战是两个领域之间的差异性。其中一个差异在于规模。不像语言翻译领域中将单词token作为基本的处理元素,视觉中的元素在规模上有很大的变化,在诸如目标检测之类的任务中,这是一个大问题。在已有的基于transformer的模型中,tokens都是固定的规模,这不适合视觉任务。另一个差异在于图像中的像素相比于文本中的单词具有更高的分辨率。很多视觉任务,如语义分割,需要在像素级别上进行密集预测,这对于transformer来说很难处理,因为其中的self-attention计算复杂度是与图像尺寸的平方相关的。为了克服这些问题,我们提出了一个通用的transformer backbone,称之为Swin Transformer,其构造了一个分层的feature maps,并具有与图像尺寸线性相关的计算复杂度。如图1(a)所示:Swin Transformer从一个小尺寸的patch开始,构造了一个分层表示,然后逐步合并相邻patch送入更深的Transformer层。有了这些层次化的feature maps,Swin Transformer可以方便地利用一些先进技术,如FPN、U-Net,进行密集预测。通过将self-attention计算限制在非重叠的窗口中,可以达到线性的计算复杂度。每个窗口中的patch数量是固定的,因此复杂度与图像大小成线性关系。这些优点,是的Swin Transformer适合作为视觉任务中的通用backbone,这与之前基于transformer的架构使用单一分辨率的feature maps且具有二次复杂度是不同的。

Swin Transformer: 使用滑动窗口的分层视觉transformer_第1张图片

 在Swin Transformer中一个关键的设计是在相邻的自注意力层间,分割窗口的“位移”,也即相邻两个层中,patch划分的窗口是不同的,如图2所示,这些移动窗口将上一层的窗口连接了起来,从而显著增强了建模能力。这种策略在实际的延迟方面也是高效的:所有的query patches共享同一个key set,这有利于内存的访问。相比之下,早期的一些基于滑窗的self-attention方法,其在普通硬件上的低延迟是由于不同的query像素具有不同key set。我们的实验表明,所提出的shifted window方法比sliding window方法具有更低的延时。

Swin Transformer: 使用滑动窗口的分层视觉transformer_第2张图片

所提出的Swin Transformer在图像分类、目标检测、语义分割等任务上达到了超强的性能,在这三个任务上显著超过了ViT/DeiT、ResNe(X)t。其可以 在COCO test-dev set达到58.7 box AP、51.1 mask AP,超过先前的SOTA方法;在ADE20K语义分割上,可以达到53.5 mIoU,超出SETR+3.2mIoU;其在ImageNet-1K 图像分类上也达到了86.4%的top-1准确率。

我们相信,一个跨计算机视觉和自然语言处理的统一架构将有利于这两个领域,因为它将促进视觉和文本信号的联合建模,并且可以更深入地共享这两个领域的建模知识。我们希望Swin Transformer在各种视觉问题上的出色表现,可以在社区中推动这种信念的深入,并鼓励视觉和语言信号的统一建模。

2、相关工作

CNN及其变体:

CNN是整个CV领域的标准网络模型。CNN已经出现几十年了,但直到AlexNet出现才成为主流。自此以后,更深、更高效的CNN网络结构被提出,进一步推进了CV领域的深度学习浪潮,如VGG、GoogleNet、ResNet、DenseNet、HRNet、EfficientNet等。除了网络架构上的改进,也有一些工作致力于研究一些特别的卷积层,如深度可分离卷积、可变性卷积。虽然CNN及其变体仍是CV应用的主流backbone,但我们强调了Transformer类的架构在视觉和语言之间统一建模的强大潜力。我们的工作在几个基本的视觉识别任务上取得了优异的成绩,希望它能够对建模方式做出一些改变。

基于self-attention的backbone:

受self-attention和transformer在NLP领域成功的激励,一些工作利用self-attention层部分或者全部替代了ResNet中的空间卷积层。在这些工作中,self-attention是在一个局部窗口中计算的,每个像素都参与优化,并且相比ResNet架构能够达到稍优的准确率/FLOPs

平衡。然而,其昂贵的访存导致了明显大于CNN的延迟。我们没有使用滑窗(sliding windows)机制,而是提出了在相邻层使用变换窗口(shift windows),这使得在一般硬件中的实现更为高效。

使用self-attention/Transformer补充CNN:

另一个研究路线是使用self-attention或者transformer来强化标准CNN。self-attention层可以通过提供编码长距离依赖或者异构交互的能力,来补足backbone或者head网络。最近,在Transformer中的Encoder-Decoder设计已被应用于目标检测和实例分割任务中。我们的工作探索了transformer对于基本视觉特征提取的适应性,属于这类工作的补充。

基于transformer的视觉backbone:

与我们工作最接近的是ViT及其后续变体。ViT开创性的在中等大小的图像patches上使用transformer,来进行图像分类。其在图像分类任务上,相比于传统CNN网络,达到了令人瞩目的速度-精度平衡。然而,ViT需要使用大规模的训练集(也即JFT-300M)才能达到不错的效果,所以DeiT引入了一些训练策略使得ViT在更小的ImageNet-1K数据集同样有效。尽管ViT在图像分类上的效果令人鼓舞,但它的架构其并不适合作为密集预测任务的通用backbone,且当输入图像分辨率较高时,其计算复杂度平方增加。也有几项工作使用ViT进行密集预测任务,如目标检测、语义分割,这些工作直接对ViT的feature maps进行上采样或者反卷积,不过其性能都比较差。与我们的工作同时进行的,还有对ViT进行修改以进行更好的图像分类。我们发现,所提出的Swin Transformer相比这些方法,能够在图像分类上达到最好的速度-精度平衡,尽管Swin Transformer更关注的是通用性能。另一项同时进行的工作探索了一个相似的思路:在transformer上构造多分辨率的feature maps。其计算复杂度仍是图像大小的二次方,而我们的是线性,且是已被证实有利于视觉建模的局部操作。我们的方法既高效又有效,在COCO目标检测和ADE20K语义分割上均达到了SOTA。

3、方法

3.1、总体架构

Swin Transformer: 使用滑动窗口的分层视觉transformer_第3张图片

Swin Transformer的整体架构如图3所示,这是一个小的版本(Swin-T)。其首先像ViT那样利用一个patch划分模块将一张输入的RGB图像分为多个patches,每个patch作为一个“token”,其特征视为原始图像的一部分。在我们的实现中,patch的大小设为4*4,因此每个patch的维度是4*4*3=48。然后使用一个线性嵌入层将这个原始值的特征投影到任意维度(记为C)。

这些patch token然后被送入几个带有修改了的self-attention运算(Swin Transformer块)的Transformer块,这些Transformer块保持了token的数量(\frac{H}{4}\times \frac{W}{4}),并与线性嵌入层一起,形成“Stage 1”。

为了生成层次化的表示,随着网络变深,通过patch合并层来减少token的数量。第一个patch合并层连接了一组2*2的相邻patches,并应用了一个4C维度的线性层。这使得token的数量减少了2*2=4(2*分辨率下采样)倍数,输出的通道维度为2C。然后使用Swin Transformer块进行特征转换,分辨率维持在\frac{H}{8}\times \frac{W}{8}上。这个第一次进行patches合并和特征转换的块命名为“Stage 2”。该过程重复两次,分别命名为“Stage 3”和“Stage 4”,输出分辨率分别为\frac{H}{8}\times \frac{W}{8}\frac{H}{16}\times \frac{W}{16}。这些stages共同产生了一个层次化的表示,与经典的CNN网络,如VGG、ResNet等网络的feature maps分辨率一样。因此,所提出的这种架构可以方便地在多种视觉任务上替换已有方法中的backbone。

Swin Transformer block:Swin Transformer构建时, 将标准的多头self-attention(MSA)替换为基于shifted Windows(于3.2节描述)的Transformer块,而其他模块保持不变。如图3(b)所示,一个Swin Transformer块由一个基于 shifted window的MSA模块、一个2层的中间带有GELU非线性激活函数的MLP组成。在每个MSA和MLP之前,还应用了一个层归一化(LN),并在每个模块之后都应用了一个残差连接。

3.2、基于 Shifted Window的self-attention

在标准的transformer及其改进版中,都执行了全局self-attention,也即一个token和所有其他token之间的关系都被计算了。这种全局计算导致了与token个数二次相关的复杂度,从而使其不适用于那些需要大量token来进行密集预测或者生成一个高分辨率图像的视觉任务。

在非重叠窗口中的self-attention:

为了高效建模,我们提出在局部窗口内执行self-attention。所述窗口均匀地、非重叠的分布在整个图像上。假定每个窗口包含了M*M个patch,全局MSA模块和基于窗口的MSA模块(一个图像分为h*w个patches)的计算复杂度分别表示为:

\begin{array}{l} \Omega(\mathrm{MSA})=4 h w C^{2}+2(h w)^{2} C \\ \Omega(\mathrm{W}-\mathrm{MSA})=4 h w C^{2}+2 M^{2} h w C \end{array}

前者是hw的二次方,后者当M固定时(默认为7)则为线性。对于一个较大的hw,全局self-attention计算通常过于昂贵,而基于窗口的则相对来说比较廉价。

在连续块之间的Shifted window划分:

基于窗口的self-attention模块缺少了跨窗口的连接,这限制了建模能力。为了在保持非重叠窗口的高效计算的同时引入跨窗口的连接,我们提出了一个shifted window划分方法,在连续的Swin Transformer块间交替进行两种划分配置。

如图2所示,第一个模块从左上角的像素开始使用了一种常规的窗口划分策略,一个8*8的feature map被分为2*2的窗口,每个窗口的大小为4*4(M=4)。然后,下一个模块使用了一种与上个模块不同的窗口配置,使用\left\lfloor\frac{M}{2}\right\rfloor,\left\lfloor\frac{M}{2}\right\rfloor的窗口替换原来的窗口,也即图2右图。

使用了这种Shifted window的划分策略后,则连续的Swin Transformer块的计算可表示为:

\begin{array}{l} \hat{\mathbf{z}}^{l}=\mathrm{W}-\operatorname{MSA}\left(\mathrm{LN}\left(\mathbf{z}^{l-1}\right)\right)+\mathbf{z}^{l-1} \\ \mathbf{z}^{l}=\operatorname{MLP}\left(\mathrm{LN}\left(\hat{\mathbf{z}}^{l}\right)\right)+\hat{\mathbf{z}}^{l} \\ \hat{\mathbf{z}}^{l+1}=\mathrm{SW}-\operatorname{MSA}\left(\mathrm{LN}\left(\mathbf{z}^{l}\right)\right)+\mathbf{z}^{l} \\ \mathbf{z}^{l+1}=\mathrm{MLP}\left(\mathrm{LN}\left(\hat{\mathbf{z}}^{l+1}\right)\right)+\hat{\mathbf{z}}^{l+1} \end{array}

其中,\hat z^lz^l分别表示(S)W-MSA和MLP模块的输出;W-MSA和SW-MSA分别表示使用常规窗口划分的MSA(也即图2左图)、使用Shifted Window划分方法的MSA(也即图2右图)。

Shifted Window划分方法对上层的相邻非重叠窗口做了连接,且发现这种做法对图像分类、目标检测、语义分割均有效。

针对shifted配置的高效批计算:

当使用shifted window进行划分时,将会导致一个问题:会产生更多窗口,从\left\lceil\frac{h}{M}\right\rceil \times\left\lceil\frac{w}{M}\right\rceil\left(\left\lceil\frac{h}{M}\right\rceil+1\right) \times \left(\left\lceil\frac{w}{M}\right\rceil+1\right),以及一些可能会小于M*M的窗口(为了使窗口能够被feature map的大小整除,可能会在右下角进行padding)。一个朴素的解决方案是:将更小的窗口padding到M*M,并在计算注意力时mask掉这些padding的值。当常规划分方法中窗口个数比较少时,比如2*2,这种朴素方法所增加的计算量是相当大的(2*2->3*3,增大了2.25倍)。在这里,我们提出了一个更高效批量计算方法,通过循环转移(cyclic shift)的方法,把原来的shifted 配置的窗口变为常规划分的窗口,如图4所示:

Swin Transformer: 使用滑动窗口的分层视觉transformer_第4张图片

经过转换后的窗口就可能由几个不相邻的feature map组成,因此后续使用masking机制将self-attention计算限制在各自的子窗口内。如此一来,窗口的个数就与常规划分方式一致了,也因此同样高效。

相对位置偏差:

在计算self-attention时,我们按照已有的做法,在每个head计算相似性过程中加入了相对位置偏差B \in \mathbb R^{M^2 \times M^2}

\text { Attention }(Q, K, V)=\operatorname{SoftMax}\left(Q K^{T} / \sqrt{d}+B\right) V

其中,Q, K, V \in \mathbb{R}^{M^{2} \times d}是query、key、value矩阵;d是query/key的维度,M^2是一个窗口中patches的个数。由于各轴上的相对位置在[−M +1,M−1]范围内,我们将一个更小尺寸的偏移矩阵\hat{B} \in \mathbb{R}^{(2 M-1) \times(2 M-1)}进行参数化,且B的值来自于\hat B

我们观察到,使用了相对位置偏置后,相比没有使用或者使用了绝对位置嵌入的项,具有显著的提升。而进一步添加绝对位置嵌入,性能会有所下降,所以我们没有采用。

在预训练阶段学习到的相对位置偏置也可以通过使用双三次插值来对不同窗口大小的模型在微调时初始化。

3.3、架构变体

我们构造的基本模型为Swin-B,其模型大小及计算复杂度与ViT-B和DeiT-B相当。我们同时也构造了Swin-T、Swin-S以及Swin-L,这些模型是模型大小及计算复杂度分别为0.25x、0.5x和2x倍的版本。注意,Swin-T和Swin-S的复杂度分别相当于ResNet-50(DeiT-S)和ResNet-101。窗口大小默认设为M=7;在所有实验中,每个head的query维度设为d=32,每个MLP的膨胀层均为\alpha=4。各模型的架构超参分别为:

  • Swin-T: C = 96, layer numbers = {2,2,6,2}
  • Swin-S: C = 96, layer numbers ={2,2,18,2}
  • Swin-B: C = 128, layer numbers ={2,2,18,2}
  • Swin-L: C = 192, layer numbers ={2,2,18,2}

其中,C是第一阶段隐藏层的通道数。模型大小、理论上的计算复杂度(FLOPs),以及在ImageNet上模型的吞吐量见表1.

4、实验结果

Swin Transformer: 使用滑动窗口的分层视觉transformer_第5张图片

Swin Transformer: 使用滑动窗口的分层视觉transformer_第6张图片

Swin Transformer: 使用滑动窗口的分层视觉transformer_第7张图片

Swin Transformer: 使用滑动窗口的分层视觉transformer_第8张图片

Swin Transformer: 使用滑动窗口的分层视觉transformer_第9张图片

5、结论

 本文提出的Swin Transformer主要有两个特点:feature map的层次化表示、线性计算复杂度。所提出的模型能够在COCO目标检测、ADE20K语义分割上达到新的SOTA。

你可能感兴趣的:(DeepLearning,论文笔记,SwinTransformer,分层表示,线性计算复杂度,视觉Transformer)