本文提出了一种新的视觉Transformer,称为Swin Transformer(Shifted WINdows Transformer),它可以作为计算机视觉的通用backbone。在视觉和文本两个领域之间的差异,比如视觉实体的规模差异很大,图像中的像素与文本中的单词相比分辨率很高,这就给从语言到视觉的转换带来了挑战。为了解决这些差异,我们提出了一种分层Transformer,其表示是用移位窗口(shifted window)计算的。移位窗口将自注意计算限制在非重叠的局部窗口上,同时允许跨窗口连接,从而提高了效率。这种分层结构具有在各种尺度下建模的灵活性,并且相对于图像大小具有线性计算复杂度。Swin Transformer的这些特性使其与广泛的视觉任务兼容,包括图像分类(ImageNet-1K上为87.3 top-1 Acc)和密集预测任务,如目标检测(COCO testdev上的 58.7 box AP和 51.1 mask AP)和语义分割(ADE20K val上为 53.5 mIoU)。它的性能大大超过了之前的技术水平,COCO上的 box AP 和 mask AP 分别为+2.7和+2.6,ADE20K上的 mIoU 为+3.2,显示了基于Transformer的模型作为视觉backbone的潜力。分层设计和移位窗口方法也被证明对 all-MLP 架构都是有益的。
计算机视觉中的建模一直由卷积神经网络(CNN)主导。从AlexNet及其在ImageNet图像分类挑战中的革命性表现开始,CNN通过扩大规模,更复杂的连接,更丰富的卷积变体变得越来越强大。CNN作为各种视觉任务的backbone网络,这些体系结构的进步导致了性能的提高,从而广泛提升了整个领域。
另一方面,自然语言处理(NLP)中网络体系结构的演变走了一条不同的道路,今天流行的体系结构是Transformer。Transformer是为sequence modeling和transduction tasks而设计的,它以关注数据中的长期依赖关系而闻名。它在语言领域的巨大成功促使研究人员研究了它对计算机视觉的适应性,最近它在某些任务上展示了有希望的结果,特别是图像分类和联合视觉-语言建模。
在本文中,我们试图扩展Transformer的适用性,使其可以像CNN在视觉中一样,作为计算机视觉的通用主干。我们观察到,将其在语言领域的高性能转移到视觉领域的challenges可以通过两种模式之间的差异来解释。其中一个差异涉及规模scale。与作为语言Transformer处理基本元素的word token不同,视觉元素在规模上可能会有很大差异,这是一个在目标检测等任务中受到关注的问题。在现有的基于Transformer的模型中,token都是固定scale的,这种属性不适合这些vision应用。另一个区别是,与文本段落中的单词相比,图像中的像素分辨率要高得多。存在许多视觉任务,例如需要在像素级进行密集预测的语义分割,而这对于高分辨率图像上的Transformer来说是很困难的,因为它自身注意力的计算复杂性是图像大小的二次方。为了克服这些问题,我们提出了一种通用的Transformer主干,称为Swin Transformer,它构造了层次特征图(hierarchical feature maps),并具有与图像大小成线性关系的计算复杂度。如图1a所示,Swin Transformer通过从小尺寸的patches(灰色网格)开始,逐渐将相邻patches合并到更深的Transformer层中来构建层次表示。有了这些分层特征映射,Swin Transformer模型可以方便地利用其他技术进行密集预测,如特征金字塔网络FPN或U-Net。线性计算复杂度是通过在分割图像的非重叠窗口(用红色标出)内局部计算自注意力来实现的。每个窗口中的patches数是固定的,因此复杂性与图像大小成线性关系。这些优点使Swin Transformer适合作为各种视觉任务的通用主干,这与以前基于Transformer的体系结构不同,后者生成单一分辨率的特征图,并且具有二次复杂性。
Swin Transformer的一个关键设计元素是在连续的自注意力层之间切换窗口分区,如图2所示。移动窗口桥接了前一层的窗口,提供了它们之间的连接,显著增强了建模能力。这种策略对于real world的计算也是有效的:一个窗口中的所有query patches都共享相同的key set,这有助于硬件中的内存访问(query和key是自我注意层中的投影向量)。相比之下,早期基于滑动窗口(sliding window)的自注意方法由于对于不同query pixel有不同的key set,在通用硬件上的延迟较低。(有一些有效的方法可以在通用硬件上实现基于sliding window的卷积层,是由于其在整个特征图中共享kernel权重,但是基于滑动窗口的自注意力层在实践中很难有效地访问内存)。我们的实验表明,所提出的移位窗口(shifted window)方法的延迟比滑动窗口方法低得多,但建模能力相似。事实证明,移位窗口方法也适用于所有MLP体系结构。
我们相信,跨越计算机视觉和自然语言处理的统一体系结构可以使这两个领域受益,因为它将促进视觉和文本信号的联合建模,并且来自这两个领域的建模知识可以更深入地共享。我们希望Swin Transformer在各种视觉问题上的出色表现能够加深社区的这种信念,并鼓励视觉和语言信号的统一建模。
滑动窗口(Sliding Windows)在目标检测过程中的作用是定位目标(物体、动物等)在图片中的位置。在计算机视觉中滑动窗口是一个矩形框,它沿着从左向右、从上向下的方向在图片上滑动以达到提取出图片中每一个区域的目的;
对于图像处理,sliding window加上权重参数其实就是卷积;
CNN是贯穿计算机视觉的标准网络模型。虽然CNN已经存在了几十年,但直到AlexNet的引入,CNN才开始发展并成为主流。此后,更深更复杂的网络被提出进一步推动了计算机视觉的发展。除了这些架构上的进步,还有很多关于改进单个卷积层的工作,比如可变形卷积。虽然CNN及其变体仍然是计算机视觉应用的主要主干架构,但我们强调了Transformer架构在视觉和语言之间统一建模方面的巨大潜力。我们的工作在几个基本的视觉识别任务上取得了很好的表现,我们希望这将有助于建模的转变。
同样受NLP领域中自注意力层和Transformer架构的启发,一些工作使用自注意力层来取代流行的ResNet中的部分或全部空间卷积层。在这些工作中,自注意是在每个图像的局部窗口内计算的,以加快优化,并且它们实现了比对应的ResNet架构稍好的精度与速度权衡。然而,昂贵的内存访问导致它们的实际延迟明显大于卷积网络。我们建议在连续层之间切换窗口(shifted window),而不是使用滑动窗口(sliding window),这允许在一般硬件中更有效地实现。
另一项工作是用自注意力层或Transformer来增强标准CNN架构。自注意力层可以补充主干网或头部网络,提供编码远程依赖或异构交互的能力。最近,Transformer中的编码器-解码器设计已应用于对象检测和实例分割任务。我们的工作探索了Transformer对基本视觉特征提取的适应性,是对这些工作的补充。
与我们的工作最相关的是视觉Transformer(ViT)及其后续变体。ViT的开创性工作直接将Transformer架构应用于非重叠的中等大小图像patches,用于图像分类。与卷积网络相比,它在图像分类上实现了令人印象深刻的速度-精度折衷。虽然ViT需要大规模的训练数据集(即JFT-300M)才能表现良好,但DeiT引入了几种训练策略,使ViT也能有效地使用较小的ImageNet-1K数据集。ViT关于图像分类的研究令人鼓舞,但是,由于其低分辨率特征映射和复杂度随图像大小的二次增加,其体系结构不适合用作密集视觉任务或输入图像分辨率较高时的通用主干网络。有几项工作将ViT模型应用于通过直接上采样或反卷积进行的目标检测和语义分割的稠密视觉任务,但性能相对较低。与我们的工作同时进行的还有一些修改ViT体系结构,以实现更好的图像分类。根据经验,我们发现我们的Swin Transformer体系结构在这些图像分类方法中实现了最佳的速度精度权衡,尽管我们的工作侧重于通用性能,而不是专门针对分类。另一项并行工作探索了在Transformer上构建多分辨率特征图的类似思路。它的复杂度仍然是图像大小的二次方,而我们的是线性的,并且在局部操作,这在建模视觉信号的高相关性方面已经证明是有益的。我们的方法既高效又有效,在COCO目标检测和ADE20K语义分割方面都达到了最先进的精度。
图3给出了Swin Transformer结构的概述,它展示了tiny版本(Swin-T)。它首先通过patches分割模块(patch splitting module)将输入的RGB图像分割为非重叠patches(张量名为Patch Partition)。每个patch被视为“token”,其特征被设置为原始像素RGB值的concatenation。在我们的实现中,我们使用 4 × 4 4\times 4 4×4的patch size,因此每个patch即token的维数为 4 × 4 × 3 = 48 4\times 4\times 3=48 4×4×3=48。在该原始值特征上应用线性embedding层,将其投影到任意维度(记为 C C C)。
在这些patch token上应用了几个自注意计算(Swin Transformer blocks)的Transformer block。Transformer block维持了token的数量 ( H 4 × W 4 ) (\frac{H}{4}\times\frac{W}{4}) (4H×4W),它与线性embedding一起被称为“Stage1”。
为了产生分层表示,随着网络的深入,通过patches合并层(patch merging layers)来减少token的数量。第一个patch合并层连接每组 2 × 2 2×2 2×2相邻patch的特征,并在 4 C 4C 4C维( 4 C = 2 × 2 × C 4C=2\times 2\times C 4C=2×2×C)连接的特征上应用线性层。这将token的数量减少了 2 × 2 = 4 2×2=4 2×2=4的倍数(分辨率的 2 × 2× 2×降采样),并且输出维度设置为 2 C 2C 2C。然后应用Swin Transformer blocks进行特征变换,分辨率保持在 H 8 × W 8 \frac{H}{8}×\frac{W}{8} 8H×8W。第一个patch merging和feature transformation被称为“Stage2”。这个过程被重复两次,分别为“Stage3,Stage4”,输出分辨率为 H 16 × W 16 \frac{H}{16}×\frac{W}{16} 16H×16W和 H 32 × W 32 \frac{H}{32}×\frac{W}{32} 32H×32W。这些Stage共同产生一个层次表示,具有与卷积网络相同的特征图分辨率,例如VGG和ResNet。因此,所提出的体系结构可以方便地替代现有方法中用于各种视觉任务的主干网络。
Swin Transformer block
Swin Transformer是通过将Transformer块中的标准多头自注意力(MSA,multi-head self attention)模块替换为基于移位窗口的模块(如3.2节所述)构建的,其他层保持不变。如图3b所示,Swin Transformer block由一个基于移位窗口的MSA模块组成,然后是一个中间带有GELU非线性的两层MLP。在每个MSA模块和每个MLP之前应用LayerNorm(LN)层,在每个模块之后应用残差连接。
残差连接是一种High way的连接方式,有利于梯度回传,Highway神经网络在输入与输出之间建立直接连接,使得输出层可以直接将梯度回传到输入层,从而避免因网络层数过多而带来的梯度爆炸或弥散问题
标准Transformer体系结构及其对图像分类的适应版本都进行了全局自注意力计算,其中计算了token和所有其他token之间的关系。全局计算导致token数量的二次方复杂度,这使得它不适用于许多需要大量token进行密集预测或表示高分辨率图像的视觉问题。
非重叠窗口中的自注意力
为了有效地建模,我们建议在局部窗口内计算自注意。窗口被布置成以不重叠的方式均匀地分割图像。假设每个窗口包含 M × M M\times M M×M个patches(图像中一共 h × w h\times w h×w 个 patches),全局MSA模块和基于窗口的计算复杂度为: Ω ( M S A ) = 4 h w C 2 + 2 ( h w ) 2 C \Omega(MSA)=4hwC^{2}+2(hw)^{2}C Ω(MSA)=4hwC2+2(hw)2C Ω ( W − M S A ) = 4 h w C 2 + 2 M 2 h w C \Omega(W-MSA)=4hwC^{2}+2M^{2}hwC Ω(W−MSA)=4hwC2+2M2hwC其中,前者是patch数量 h w hw hw的二次方,后者是线性的(当 M M M固定时)。对于很大的 h w hw hw来说,全局自注意计算通常是负担不起的,而基于窗口的自注意是可伸缩的。
两个复杂度计算没有考虑SoftMax;
回顾矩阵乘法的复杂度
对于矩阵 A n , m , B m , n A_{n,m},B_{m,n} An,m,Bm,n,如果 C = A B C=AB C=AB,则复杂度有 O ( n m n ) = O ( n 2 m ) O(nmn)=O(n^{2}m) O(nmn)=O(n2m),原因如下:
for(i=0;i<n;i++){ //A矩阵中的n
for(j=0;j<m;j++){ //A矩阵中的m
for(k=0;k<n;k++){ //B矩阵中的n
C[i][j]= C[i][j]+A[i][k]*B[k][j];}}}
对于三矩阵乘法 A n , m , B m , n , C n , m A_{n,m},B_{m,n},C_{n,m} An,m,Bm,n,Cn,m,由于 A B = D n , n = O ( n 2 m ) AB=D_{n,n}=O(n^{2}m) AB=Dn,n=O(n2m),则 E = A B C = D n , n C n , m = O ( n n m ) = O ( n 2 m ) E=ABC=D_{n,n}C_{n,m}=O(nnm)=O(n^{2}m) E=ABC=Dn,nCn,m=O(nnm)=O(n2m),在实际运行时,矩阵乘法是二元运算的,所以 A B C ABC ABC会分开执行,因此总的复杂度为 O ( n 2 m ) + O ( n 2 m ) = O ( n 2 m ) O(n^{2}m)+O(n^{2}m)=O(n^{2}m) O(n2m)+O(n2m)=O(n2m)。
对于MSA的复杂度:
对于W-MSA:W-MSA与MSA的计算过程一致,区别在于W-MSA在窗口内计算,且每个窗口有 M × M M\times M M×M个patches,图像有 h M × w M \frac{h}{M}\times\frac{w}{M} Mh×Mw个窗口要计算。所以它的复杂度为: ( h M × w M ) × ( 4 M 2 C 2 + 2 M 4 C ) = 4 h w C 2 + 2 M 2 h w C (\frac{h}{M}\times\frac{w}{M})\times(4M^{2}C^{2}+2M^{4}C)=4hwC^{2}+2M^{2}hwC (Mh×Mw)×(4M2C2+2M4C)=4hwC2+2M2hwC
连续Swin Transformer blocks中的移位窗口划分
基于窗口的自注意力模块缺乏跨窗口连接,这限制了其建模能力。为了在保持非重叠窗口高效计算的同时引入跨窗口连接,我们提出了一种移位窗口划分方法,该方法在连续的Swin Transformer block中交替使用两种划分配置。
如图2,第一个模块使用常规窗口划分策略(从左上角像素开始),将 ( 8 , 8 ) (8,8) (8,8)特征图划分为 2 × 2 2\times 2 2×2个大小为 ( 4 , 4 ) (4,4) (4,4)的窗口( M = 4 M=4 M=4)。然后,下一个模块通过将窗口从规则分区的窗口起点移动到 ( M 2 , M 2 ) (\frac{M}{2},\frac{M}{2}) (2M,2M)像素再进行划分,即采用与前一层的窗口配置不同的窗口配置。
第一个block使用常规窗口划分策略,第二个block使用移动窗口策略,window从特征图的 ( M 2 , M 2 ) (\frac{M}{2},\frac{M}{2}) (2M,2M)开始,然后进行窗口划分。
使用移位窗口划分方法,连续的Swin Transformer blocks计算如下: z ^ l = W − M S A ( L N ( z l − 1 ) ) + z l − 1 \widehat{z}^{l}=W-MSA(LN(z^{l-1}))+z^{l-1} z l=W−MSA(LN(zl−1))+zl−1 z l = M L P ( L N ( z ^ l ) ) + z ^ l z^{l}=MLP(LN(\widehat{z}^{l}))+\widehat{z}^{l} zl=MLP(LN(z l))+z l z ^ l + 1 = S W − M S A ( L N ( z l ) ) + z l \widehat{z}^{l+1}=SW-MSA(LN(z^{l}))+z^{l} z l+1=SW−MSA(LN(zl))+zl z l + 1 = M L P ( L N ( z ^ l + 1 ) ) + z ^ l + 1 z^{l+1}=MLP(LN(\widehat{z}^{l+1}))+\widehat{z}^{l+1} zl+1=MLP(LN(z l+1))+z l+1其中, z ^ l \widehat{z}^{l} z l和 z l z^{l} zl表示(S)W-MSA和MLP的输出特征(对于block l l l)。W-MSA和SW-MSA分别表示使用规则窗口划分和移位窗口划分策略的基于窗口的多头自注意力。
移位窗口分割方法引入了前一层中相邻非重叠窗口之间的连接,并被发现在图像分类、对象检测和语义分割方面是有效的。
移位策略的高效批量计算
shifted window分区的一个问题是,会导致出现更多的窗口,从 ( h M , w M ) (\frac{h}{M},\frac{w}{M}) (Mh,Mw)个变成 ( h M + 1 , w M + 1 ) (\frac{h}{M}+1,\frac{w}{M}+1) (Mh+1,Mw+1),并且有些窗口大小是小于 M × M M\times M M×M的(比如图2)。一个简单的解决方案是将大小不足 M × M M\times M M×M的窗口填充到 M × M M×M M×M的大小,并在计算注意力时屏蔽这些填充值(相当于0填充)。但是这种navie的做法会增加很多计算量(比如计算的窗口数量从 2 × 2 2\times 2 2×2变成 3 × 3 3\times 3 3×3)。这种navie的做法其实就是sliding window。在这里,我们提出了一种更有效的批量计算方法,方法是向左上方向循环移位(cyclic shift),如图4所示。在这个移位之后,一个批处理窗口可能由几个在特征图中不相邻的子窗口组成,因此使用mask机制将自注意计算限制在每个子窗口内。通过循环移位,批处理窗口的数量与常规窗口分区的数量相同,因此也是有效的。
torch.roll
实现,第一次将第一行移动到最后一行,第二次将第一列移动到最后一列),从而使得最后的feature map依然为 2 × 2 2\times 2 2×2 的windows,保持原有的计算量,然后再使用图中紫色部分的masked MSA进行计算。结束之后,再reverse cyclic shift。注意masked MSA,当我们做cyclic shift后有:
对于第二个特征图,我们可以很几何地按照之前的window划分方式(标准的 2 × 2 2\times 2 2×2个窗口)去计算,但是对于3个窗口:即 (4+6),(2+8),(1+3+7+9) 的attention会混在一起,所以我们需要在计算每个窗口时进行mask MSA。
以窗口 (4+6) 为例,假设该窗口一共有4个patch:
当自注意力计算时,重点在于 Q K T QK^{T} QKT,为了保证信息只在cyclic shift前的window内交互,我们要确保只存在属于 window 4和 window 4的两个patch计算attention,换言之就是在计算注意力时候只有行和列属于相同编号的元素才保留,其他元素都mask。如下图所示:
至此,我们利用shift后的feature,和上面说的mask结合,就能得到正确的MSA结果。我们最后把shift还原,即reverse shift。
相对位置偏差
在计算自注意力时,我们计算相对位置偏差 B ∈ R M 2 × M 2 B\in R^{M^{2}\times M^{2}} B∈RM2×M2,对每个注意力head计算: A t t e n t i o n ( Q , K , V ) = S o f t M a x ( Q K T d + B ) V Attention(Q,K,V)=SoftMax(\frac{QK^{T}}{\sqrt{d}}+B)V Attention(Q,K,V)=SoftMax(dQKT+B)V其中, Q , K , V ∈ M M 2 × d Q,K,V\in M^{M^{2}\times d} Q,K,V∈MM2×d分别是query,key,value矩阵。 d d d是维数, M 2 M^{2} M2是每个窗口的patches数。
我们观察到与没有这个偏差项或使用绝对位置embedding的对应项相比,有显著的改进。进一步向输入中添加绝对位置embedding会略微降低性能,因此我们的实现中不采用绝对位置embedding。
预训练中学习到的相对位置偏差也可用于初始化模型。
我们建立了一个名为Swin-B的基础模型,该模型的大小和计算复杂度与ViTB/DeiT-B相似。我们也引入了Swin-T,Swin-S和Swin-L,模型大小和计算复杂度分别是基础模型的0.25,0.5和2倍。注意,Swin-T和Swin-S的复杂性分别与ResNet-50(DeiT-S)和ResNet-101的复杂性相似。窗口大小默认为 M = 7 M=7 M=7,每个head的维数为 d = 32 d=32 d=32,每个MLP的层数为 α = 4 \alpha=4 α=4,对于所有的实验。这些变体的架构参数包括:
其中, C C C是第一个stage的hidden layer的通道数,ImageNet图像分类模型变体的模型大小、理论计算复杂度(FLOPs)和吞吐量(throughput)如表1所示。
我们在ImageNet-1K图像分类、COCO目标检测和ADE20K语义分割上进行了实验。在下文中,我们首先比较了所提出的Swin Transformer架构与之前在这三项任务上的最新进展。然后,我们讨论了SwinTransformer的重要设计元素。
Settings
对于图像分类,我们在ImageNet-1K上对Swin Transformer进行了基准测试,其中包含1.28M个训练图像和来自1000个类的50K个验证图像。报告了单个类别的最高精度。我们考虑两种训练环境:
Results with regular ImageNet-1K training
表1a给出了使用常规ImageNet-1K训练与其他主干网(包括基于Transformer和基于ConvNet)的比较。
与之前最先进的基于Transformer的架构(即DeiT)相比,SwinTransformer明显超过了具有类似复杂性的对应DeiT架构:使用224输入的Swin-T(81.3%)比DeiT-S(79.8%)高出1.5%,使用224或384输入的Swin-B(83.3%或84.5%)比DeiT-B(81.8%或83.1%)高出+1.5%或1.4%。
与最先进的ConvNets(即RegNet)相比,SwinTransformer实现了略好的速度精度折衷。
Results with ImageNet-22K pre-training
我们还在ImageNet22K上预训练更大容量的Swin-B和Swin-L。ImageNet-1K图像分类的微调结果如表1b所示。对于Swin-B,ImageNet22K预训练带与ImageNet-1K从头开始的训练相比,收益为1.9%。与之前ImageNet-22K预训练的最佳结果相比,我们的模型实现了显著更好的速度-精度权衡:Swin-B获得86.4%的top-1精度,比具有类似推理吞吐量的ViT高2.4%,并且FLOPs更低。更大的Swin-L模型达到87.3%的top-1精度,比Swin-B模型高+0.9%。
Settings
目标检测和实例分割实验在COCO 2017上进行,其包含了118K训练图像,5K验证图像,20K测试图像(test-dev)。使用验证集进行Ablation study,并在test-dev上比较。对于Ablation study,我们考虑四种典型的目标检测框架:Mask R-CNN、ATSS、RepPoints v2和Sparse RCNN。对于这四个框架,我们使用相同的设置:多尺度训练(调整输入大小,使短边在480到800之间,而长边最多为1333),AdamW优化器(初始学习率为0.0001,weight decay为0.05,batch大小为16),以及3xschedule(36个epoch)。另外,我们采用Soft-NMS和ImageNet-22K预训练模型。
Setting
ADE20K是一个广泛使用的语义分割数据集,涵盖了150个语义类别。它总共有25K个图像,其中20K用于培训,2K用于验证,另外3K用于测试。我们利用UperNet作为我们的基础框架,以实现其高效性。
1.时间复杂度的减少,主要是在于分别只在window内计算自注意力;
2.为了融合不同窗口之间的信息,采用shifted window划分策略,不同于过去的sliding window这种navie的方式,shifted window的特点在于cyclic shift + masked MSA + reverse cyclic shift,这实现了无padding且不增加窗口数量的情况下达到sliding window的效果;
3.层级的架构可以考虑不同尺度的Window,从而获得多尺度信息;
4.可扩展性强,是更好的视觉Transformer,对下游任务均有提升。