标题:Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
会议:ICCV2021
论文地址:https://ieeexplore.ieee.org/document/9710580/
官方代码:https://github.com/microsoft/Swin-Transformer
作者单位:微软亚洲研究院、中国科学技术大学、西安交通大学、清华大学
文章目录
- Abstract
- 1. Introduction
- 2. Related Work
- 3. Method
- 3.1. Overall Architecture
- 3.2. Shifted Window based Self-Attention
- 3.3. Architecture Variants
- 4. Experiments
- 4.1. Image Classification on ImageNet-1K
- 4.2. Object Detection on COCO
- 4.3. Semantic Segmentation on ADE20K
- 4.4. Ablation Study
- 5. Conclusion
Abstract
本文提出了一种新的视觉Transformer,称为Swin Transformer,它可以作为计算机视觉的通用主干。将Transformer从语言转换到视觉的挑战来自于两个领域之间的差异,例如视觉实体的尺度变化较大、图像中的像素相对于文本中的文字具有较高的分辨率。为了解决这些差异,我们提出了一个层级式Transformer,其表征是通过移动窗口(Shifted windows)来计算的。移动窗口方案通过将自注意力的计算限制在非重叠的局部的窗口,从而带来更高的效率,同时还允许跨窗口的联系。这种层级式的结构具有在不同尺度下建模的灵活性,并且具有与图像尺寸相关的线性的计算复杂度。Swin Transformer的这些特性使其可以兼容广泛的视觉任务,包括图像分类(ImageNet-1K的top-1准确率为87.3)和稠密预测任务,如目标检测(COCO test-dev上的58.7 box AP和51.1 mask AP)和语义分割(ADE20K val上的53.5 mIoU)。其性能在各个视觉任务基准上都超过了之前的SOTA,展示了基于Transformer的模型作为视觉主干的潜力。层级式的设计和移动窗口方法也被证明对all-MLP结构是有益的。
1. Introduction
计算机视觉中的建模长期以来由卷积神经网络(CNNs)主导。从AlexNet及其在ImageNet图像分类挑战上的革命性表现开始,CNN结构通过更大的规模、更广泛的连接和更复杂的卷积形式变得越来越强大。随着CNN作为各种视觉任务的主干网络,这些结构的进步导致了性能的提高从而广泛地改善了整个领域。
另一方面,自然语言处理(NLP)中网络结构的演化已经走向了一条不同的道路,现在盛行的结构是Transformer。为了序列建模和转换任务而设计,Transformer的特点是利用注意力来建模数据中的远程依赖关系。它在语言领域的巨大成功使研究人员开始研究它对计算机视觉的适应性,最近它在某些任务上表现出了很好的效果,特别是图像分类和联合视觉-语言建模。
在本文中,我们寻求扩展Transformer的适用性,使其可以作为计算机视觉的通用主干,就像它在NLP中和CNNs在视觉中一样。我们观察到,将其在语言领域的高性能迁移到视觉领域的重大挑战可以通过两种模式之间的差异来解释。差异之一是尺度。与语言Transformers中作为处理的基本元素的单词token不同,视觉元素可以在尺度上有很大的变化,这个问题在目标检测等任务中受到关注。在现有的基于Transformer的模型中,token都是固定尺度的,这一特性不适合这些视觉应用。另一个不同之处在于,图像中像素的分辨率要远高于文本段落中的文字。存在许多视觉任务,如语义分割,需要在像素级进行稠密预测,这对于高分辨率图像上的Transformer来说是很难解决的,因为它的自注意力的计算复杂度是图像尺寸的二次方。为了克服这些问题,我们提出了一个通用的Transformer主干,称为Swin Transformer,它构造了层级式的特征图,并且具有与图像尺寸呈线性的计算复杂度。如图1a所示,Swin Transformer从小尺寸的patch(用灰色勾勒)开始,在更深的Transformer层中逐步合并相邻的patch,从而构建层级式的表征。通过这些层级式的特征图,Swin Transformer模型可以方便地利用先进的稠密预测技术,如特征金字塔网络(FPN)或U-Net。线性的计算复杂度是通过在划分图像的非重叠窗口(用红色勾勒)内局部地计算自注意力来实现的。由于每个窗口中的patch的数量是固定的,因此复杂度与图像尺寸成线性关系。这些优点使得Swin Transformer适合作为各种视觉任务的通用主干,这与以前基于Transformer的结构不同,后者只生成单一分辨率的特征图,并且具有平方级复杂度。
Swin Transformer的一个关键设计要素是其窗口划分在连续的自注意力层之间移动,如图2所示。移动的窗口桥接了前一层的窗口,提供了它们之间的联系,显著增强了建模能力(见表4)。该策略在实际延迟方面也是有效的:窗口内的所有query块共享相同的key集合(query和key是自注意力层中的投影向量),这有利于硬件中的内存访问。相比之下,早期的基于滑动窗口的自注意力方法在通用硬件上由于不同query像素的key集合不同而导致高延迟(虽然有一些有效的方法可以在通用硬件上实现基于滑动窗口的卷积层,但由于其在特征映射中共享kernel权重,基于滑动窗口的自注意力层在实践中很难实现高效的内存访问)。我们的实验表明,所提出的移动窗口方法比之前的滑动窗口方法具有更低的延迟,但在建模能力方面是相似的(见表5和表6)。移动窗口方法也被证明对all-MLP结构是有益的。
提出的Swin Transformer在图像分类、目标检测和语义分割等识别任务上取得了强大的性能。在3个任务上,它都显著优于ViT/DeiT和ResNe(X)t模型,并且具有相似的延迟。其在COCO test-dev集上的58.7 box AP和51.1 mask AP超过了之前SOTA的结果+2.7 box AP(无外部数据的复制-粘贴)和+2.6 mask AP(DetectoRS)。在ADE20K语义分割上,它在验证集上获得了53.5 mIoU,比之前的SOTA(SETR)提高了+3.2 mIoU。在ImageNet-1K图像分类上也取得了87.3%的最高准确率。
我们相信,一个跨越计算机视觉和自然语言处理的统一结构可以使这两个领域受益,因为它将促进视觉和文本信号的联合建模,并且来自两个领域的建模知识可以更深入地共享。我们希望Swin Transformer在各种视觉问题上的强大表现能够推动这一观点在社区中深入,并鼓励对视觉和语言信号进行统一建模。
2. Related Work
CNN及其变体 卷积神经网络用作贯穿计算机视觉的标准网络模型。虽然CNN已经存在了几十年,但直到AlexNet的提出,CNN才脱颖而出成为主流。此后,更深、更有效的卷积神经网络结构被提出,以进一步推动计算机视觉领域的深度学习浪潮,如VGG、GoogleNet 、ResNet、DenseNet、HRNet和EfficientNet等。除了这些结构上的进步,在改进单个卷积层方面也有很多工作,比如深度可分离卷积和可变形卷积。虽然CNN及其变体仍然是计算机视觉应用的首要主干结构,但我们突出了Transformer式结构在视觉和语言统一建模方面的强大潜力。我们的工作在几个基本的视觉识别任务上取得了很好的表现,我们希望这将有助于建模的转变。
基于自注意力的主干结构 同样,受到自注意力层和Transformer结构在NLP领域成功的启发,一些工作使用自注意力层来代替流行的ResNet中的部分或全部空间卷积层。在这些工作中,自注意力在每个像素的局部窗口内被计算以加速优化,并且它们实现了比相应的ResNet结构稍好的准确率/FLOPs权衡。然而,它们昂贵的内存访问导致它们的实际延迟显著大于卷积网络。与之前使用的滑动窗口不同,我们提出在连续的层之间移动窗口,这允许了在通用硬件中更高效的实现。
自注意力/Transformers来补足CNN 另一个工作路线是用自注意力层或Transformers来增强一个标准的CNN结构。自注意力层可以通过提供编码远程依赖或异构交互的能力来补足主干网络或头部网络。最近,Transformer中的编码器-解码器设计已经应用于目标检测和实例分割任务。我们的工作探索了Transformers对基本视觉特征提取的适应性,是对这些工作的补充。
基于Transformer的视觉主干 与我们工作最相关的是Vision Transformer(ViT)及其后续工作。ViT的开创性工作是直接在不重叠的中等大小的图像块上应用Transformer结构进行图像分类。与卷积网络相比,它在图像分类方面取得了令人印象深刻的速度-精度权衡。虽然ViT需要大规模的训练数据集(即JFT-300M)才能表现良好,但DeiT引入了几种训练策略,使得ViT在使用较小的ImageNet-1K数据集时也是有效的。ViT在图像分类上的结果令人鼓舞,但由于其低分辨率特征图和复杂度随图像尺寸的呈二次方增加,因此其结构不适合用作稠密视觉任务或输入图像分辨率较高时的通用主干网络。有一些工作通过直接上采样或反卷积将ViT模型应用于目标检测和语义分割的稠密视觉任务,但性能相对较低。与我们的工作同时发生的是一些修改ViT结构以更好地进行图像分类的工作。根据经验,尽管我们的工作侧重于通用性能,而不是专门针对分类,但是我们发现我们的Swin Transformer结构在这些图像分类方法中实现了最佳的速度-精度权衡。另一项同时发生的工作探索了类似的思路,在Transformers上构建多分辨率特征图。它的复杂度仍然是图像大小的二次方,而我们的复杂度是线性的,并且在局部操作,这已被证明有利于建模视觉信号中的高相关性。我们的方法既高效又有效,在COCO目标检测和ADE20K语义分割上都达到了SOTA的精度。
3. Method
3.1. Overall Architecture
Swin Transformer结构的概述如图3所示,其阐明了tiny版本(Swin-T)。它首先通过像ViT一样的patch分割模块将输入的RGB图像分割成不重叠的patch。每个patch被视为一个“token”,其特征被设置为原始像素RGB值的concatenation。在我们的实现中,我们使用
4
×
4
4×4
4×4的patch大小,因此每个patch的特征维度为
4
×
4
×
3
=
48
4×4×3=48
4×4×3=48。在此raw-valued特征上使用线性嵌入(linear embedding)层,将其投影到任意维度(记为
C
C
C)。
在这些patch token上应用多个经过改进的自注意力计算(Swin Transformer块)的Transformer块。Transformer块保持token数量(
H
4
×
W
4
\frac{H}{4}×\frac{W}{4}
4H×4W)不变,与线性嵌入(linear embedding)合称为“Stage 1”。
为了产生层级式的表征,随着网络的加深,通过patch合并(patch merging)层来减少token的数量。第一个patch合并层将每组
2
×
2
2×2
2×2的相邻补丁的特征concatenate起来,并在
4
C
4C
4C维度的concatenate后的特征上使用一个线性层。这样使得token数量减少了
2
×
2
=
4
2×2=4
2×2=4(
2
×
2×
2×分辨率下采样)倍,输出维度设置为
2
C
2C
2C。之后,使用Swin Transformer块进行特征变换,分辨率保持在
H
8
×
W
8
\frac{H}{8}×\frac{W}{8}
8H×8W不变。这里的第一个patch合并与特征变换的块被记为“Stage 2”。这个过程会重复两次,也就是“Stage 3”和“Stage 4”,输出分辨率分别为
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块 Swin Transformer是通过将Transformer模块中标准的多头自注意力(multi-head self attention,MSA)模块替换为基于移动窗口的模块来构建的,其它层保持不变。如图3b所示,一个Swin Transformer模块由一个基于移动窗口的MSA模块组成,随后是一个2层的MLP,GELU非线性介于MLP之间。在每个MSA模块和每个MLP之前会应用一个LayerNorm(LN)层,在每个模块之后会应用一个残差连接。
3.2. Shifted Window based Self-Attention
标准的Transformer结构及其对图像分类的改进都是进行全局自注意力,计算一个token与所有其它token之间的关系。全局计算导致复杂度与token数量的二次方相关,使其不适用于许多需要大量token来进行稠密预测或表征高分辨率图像的视觉问题。
非重叠窗口中的自注意力 为了高效建模,我们提出在局部窗口内计算自注意力。窗口被排列成以非重叠的方式平均地划分图像。假设每个窗口包含
M
×
M
M×M
M×M个patch,全局MSA模块和在
h
×
w
h×w
h×w个patch的图像上基于窗口的计算复杂度分别为(在确定复杂度时,我们省略了SoftMax计算):
Ω
(
MSA
)
=
4
h
w
C
2
+
2
(
h
w
)
2
C
Ω
(
W-MSA
)
=
4
h
w
C
2
+
2
M
2
h
w
C
\begin{align} &\Omega(\textrm{MSA})=4hwC^2+2(hw)^2C\\ &\Omega(\textrm{W-MSA})=4hwC^2+2M^2hwC \end{align}
Ω(MSA)=4hwC2+2(hw)2CΩ(W-MSA)=4hwC2+2M2hwC其中,前者与patch数
h
w
hw
hw呈二次方关系,后者在
M
M
M固定(默认设置为7)时呈线性关系。全局自注意力计算在
h
w
hw
hw较大时通常是代价高昂的,而基于窗口的自注意力是可扩展的。
在连续的块中移动窗口划分 基于窗口的自注意力模块缺少跨窗口的联系,限制了其建模能力。为了在引入跨窗口联系的同时保持非重叠窗口的高效计算,我们提出了一种在连续的Swin Transformer块中交替使用两种划分布局的移动窗口划分方法。
如图2所示,第一个模块使用从左上角像素开始的规则的窗口划分策略,将
8
×
8
8×8
8×8的特征图均匀划分为
2
×
2
2×2
2×2个尺寸为
4
×
4
4×4
4×4(
M
=
4
M=4
M=4)的窗口。然后,下一个模块采用从前一层移动而来的窗口布局,通过把窗口从规则划分的窗口移动
(
⌊
M
2
⌋
,
⌊
M
2
⌋
)
(\left\lfloor\frac{M}{2}\right\rfloor,\left\lfloor\frac{M}{2}\right\rfloor)
(⌊2M⌋,⌊2M⌋)个像素。
采用移动窗口划分方法,连续的Swin Transformer块计算如下:
z
^
l
=
W-MSA
(
LN
(
z
l
−
1
)
)
+
z
l
−
1
z
l
=
MLP
(
LN
(
z
^
l
)
)
+
z
^
l
z
^
l
+
1
=
SW-MSA
(
LN
(
z
l
)
)
+
z
l
z
l
+
1
=
MLP
(
LN
(
z
^
l
+
1
)
)
+
z
^
l
+
1
\begin{align} &\hat{\mathrm{z}}^l=\textrm{W-MSA}(\textrm{LN}(\mathrm{z}^{l-1}))+\mathrm{z}^{l-1}\nonumber\\ &\mathrm{z}^l=\textrm{MLP}(\textrm{LN}(\hat{\mathrm{z}}^l))+\hat{\mathrm{z}}^l\nonumber\\ &\hat{\mathrm{z}}^{l+1}=\textrm{SW-MSA}(\textrm{LN}(\mathrm{z}^l))+\mathrm{z}^l\nonumber\\ &\mathrm{z}^{l+1}=\textrm{MLP}(\textrm{LN}(\hat{\mathrm{z}}^{l+1}))+\hat{\mathrm{z}}^{l+1} \end{align}
z^l=W-MSA(LN(zl−1))+zl−1zl=MLP(LN(z^l))+z^lz^l+1=SW-MSA(LN(zl))+zlzl+1=MLP(LN(z^l+1))+z^l+1其中,
z
^
l
\hat{\mathrm{z}}^l
z^l和
z
l
\mathrm{z}^l
zl别表示块
l
l
l中的
(S)W-MSA
\textrm{(S)W-MSA}
(S)W-MSA模块和
MLP
\textrm{MLP}
MLP模块的输出特征;
W-MSA
\textrm{W-MSA}
W-MSA和
SW-MSA
\textrm{SW-MSA}
SW-MSA分别表示使用规则布局和移动窗口划分布局的基于窗口的多头自注意力。
移动窗口划分方法在前一层中引入了相邻非重叠窗口之间的联系,并且在图像分类、目标检测和语义分割中被发现是有效的,如表4所示。
针对移动布局的高效批次计算 移动窗口划分的一个问题是,在移动布局中会产生更多的窗口,从
⌈
h
M
⌉
×
⌈
w
M
⌉
\left\lceil\frac{h}{M}\right\rceil×\left\lceil\frac{w}{M}\right\rceil
⌈Mh⌉×⌈Mw⌉到
(
⌈
h
M
⌉
+
1
)
×
(
⌈
w
M
⌉
+
1
)
(\left\lceil\frac{h}{M}\right\rceil+1)×(\left\lceil\frac{w}{M}\right\rceil+1)
(⌈Mh⌉+1)×(⌈Mw⌉+1),并且其中一些窗口会小于
M
×
M
M×M
M×M(为了使窗口尺寸
(
M
,
M
)
(M,M)
(M,M)可以被
(
h
,
w
)
(h,w)
(h,w)的特征图尺寸所分割,需要时会在特征图上使用右下角填充)。一个朴素的解决方案是将较小的窗口填充到
M
×
M
M×M
M×M的尺寸,并在计算注意力时掩盖填充值。当规则划分中的窗口数目较少时,例如
2
×
2
2×2
2×2时,这种朴素的解决方案增加的计算量相当大(
2
×
2
→
3
×
3
2×2→3×3
2×2→3×3,增加了2.25倍)。在这里,我们提出了一种更高效的向左上方循环移动的批次计算方法,如图4所示。在这个移动之后,一个成批的窗口可能由几个在特征图中不相邻的子窗口组成,因此使用一个掩码机制将自注意力计算限制在每个子窗口内。通过循环移动,成批的窗口的数量保持与规则窗口划分相同,因此也是高效的。该方法的低延迟如表5所示。
相对位置偏置 在计算自注意力时,我们遵循此前的工作,在计算相似度时给每个头加上一个相对位置偏置
B
∈
R
M
2
×
M
2
B\in\mathbb{R}^{M^2×M^2}
B∈RM2×M2:
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
(4)
\mathrm{Attention}(Q,K,V)=\mathrm{SoftMax}(QK^T/\sqrt d+B)V\tag{4}
Attention(Q,K,V)=SoftMax(QKT/d+B)V(4)其中,
Q
,
K
,
V
∈
R
M
2
×
d
Q,K,V\in\mathbb{R}^{M^2×d}
Q,K,V∈RM2×d分别是query、key和value矩阵;
d
d
d是query/key的维度,
M
2
M^2
M2是窗口内的patch数。由于沿每个轴的相对位置在
[
−
M
+
1
,
M
−
1
]
[-M+1,M-1]
[−M+1,M−1]范围内,因此我们参数化了一个较小尺寸的偏置矩阵
B
^
∈
R
(
2
M
−
1
)
×
(
2
M
−
1
)
\hat{B}\in\mathbb{R}^{(2M-1)×(2M-1)}
B^∈R(2M−1)×(2M−1),并且
B
B
B中的值取自
B
^
\hat{B}
B^。
我们观察到,这相对于没有该偏置项或使用绝对位置嵌入的对应方法有显著的改进,如表4所示。进一步在输入中加入绝对位置嵌入,性能会略有下降,因此在我们的实现中不采用。
在预训练中学习到的相对位置偏置也可以用来初始化模型,通过双三次插值以不同的窗口大小进行微调。
3.3. Architecture Variants
我们构建了我们的基础模型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。对于所有实验,每个头的query维度 d = 32 d=32 d=32,每个MLP的扩充层中 α = 4 \alpha=4 α=4。这些模型变体结构的超参数为:
- Swin-T: C = 96 C=96 C=96,层数 = { 2 , 2 , 6 , 2 } =\{2,2,6,2\} ={2,2,6,2}
- Swin-S: C = 96 C=96 C=96,层数 = { 2 , 2 , 18 , 2 } =\{2,2,18,2\} ={2,2,18,2}
- Swin-B: C = 128 C=128 C=128,层数 = { 2 , 2 , 18 , 2 } =\{2,2,18,2\} ={2,2,18,2}
- Swin-L: C = 192 C=192 C=192,层数 = { 2 , 2 , 18 , 2 } =\{2,2,18,2\} ={2,2,18,2}
其中, C C C是第一个阶段中隐藏层的通道数。用于ImageNet图像分类的模型变量的模型大小、理论计算复杂度(FLOPs)和吞吐量列于表1。
4. Experiments
4.1. Image Classification on ImageNet-1K
有两种训练设置:直接在ImageNet-1K训练和先在ImageNet-22K预训练然后再在ImageNet-1K微调。
具体的实验结论可以参照原文。
4.2. Object Detection on COCO
在目标检测上的对比实验,具体的实验结论可以参照原文。
4.3. Semantic Segmentation on ADE20K
在语义分割上的对比实验,具体的实验结论可以参照原文。
4.4. Ablation Study
移动窗口和相对位置偏置的有效性,具体的实验结论可以参照原文。
提出的自注意力方法的有效性,具体的实验结论可以参照原文。
5. Conclusion
本文提出了一种新的视觉Transformer,Swin Transformer,它产生一种层级式的特征表示,并且具有与输入图像尺寸相关的线性计算复杂度。Swin Transformer在COCO目标检测和ADE20K语义分割上达到了SOTA的性能,大大超过了此前最好的方法。我们希望Swin Transformer在各种视觉问题上的强大性能可以鼓励视觉和语言信号的统一建模。