打败VIT？Swin Transformer是怎么做到的

https://mp.weixin.qq.com/s/C5ZDYKPdHazR2bR9I9KFjQ

在之前的文章中，我们介绍过VIT（Vision Transformer） ，它将NLP中常用的Transformer架构用于图像分类的预训练任务上，取得了比肩ResNet的效果，成功证明了Transformer在NLP和CV上的大一统能力，进而成为后续许多工作的骨架模型。

今天我们要介绍的Swin Transformer（Shifted Window Transformer） 和VIT一样，也是一个采用Transformer架构做图像预训练的模型，但得益于它引入了CNN的归纳偏置，使得其在各类数据集上（尤其是分割、检测这样的预测密集型任务上）的效果都要优于VIT。但效果上有所提升的背后代价是：它对NLP和CV任务的统一能力下降了。这一点我们会在正文中细说。

本文在写作时，假设大家已经具备VIT相关知识。如果对VIT原理有不了解的朋友，可以参考这篇文章。

【全文目录如下】

一、Swin Transformer的诞生背景
1.1 VIT的缺陷
1.2 Swin Transformer的改进

二、Swin Transformer的整体架构

三、Patch Merging

四、W-MSA与SW-MSA
4.1 W-MSA
4.2 朴素SW-MSA
4.3 环状SW-MSA
4.4 Masked Attention
4.5 复习Swin Transformer Block

五、窗口attention计算量分析
5.1 矩阵相乘的计算量
5.2 全局attention计算量
5.3 窗口attention计算量
5.4 全局MSA与窗口MSA

六、参考

【绘图与码字不易，点赞和在看，也是持续更新的动力❤️～】

一、Swin Transformer的诞生背景

1.1 VIT的缺陷

之前在介绍VIT原理时，我们提过VIT的一个重要意义是：证明Transformer对CV和NLP的大一统性。因此VIT几乎是将Transformer encoder部分完全搬运过来（也可理解为和Bert几乎一致），然后将图片分割成pacth的形式，每个patch即等同于NLP中的一个token向量，如此一来完全以训练语料的方式做图片分类的预训练，整个过程如下：

而这样的极致统一，对于CV任务来说，有两个显著的缺陷：

• 同一实体的尺寸变化问题。 例如一张街景图，其中有很多辆车，每辆车的大小各不相同，而我们的目标是要把这些大小不同的车都检测出来。在CNN架构中，我们会对输入数据做尺寸变化，在模型的每一层输出不同大小的特征图（例如常见的UNet架构），以此教会模型探查出不同大小的实体。但是VIT由于追求和NLP任务保持一致性，它每一层的patch数量和patch大小都保持不变。这就使得VIT在分割、检测这样预测密集型的任务上注定存在弱势。

• 高分辨率图像引起的计算复杂度上升问题。在VIT中，尽管我们通过做patch的方式，减少了输入序列的长度。但VIT的attention计算仍是全局的，它的计算复杂度和图片尺寸大小呈平方关系。这也意味着VIT对高分辨率图像的处理是昂贵的。在后文中，我们会对这个复杂度的计算做更详细的说明。

1.2 Swin Transformer的改进

针对上述两点关于VIT的缺陷，Swin Transformer提出了如下改进：

图中演示的是不同layer的情况。其中，红色的方框表示窗口（local window），灰色的方框表示patch。我们规定只在窗口的范围内计算attention。

我们先来看图（b），它表示VIT的架构，具有以下特点：

• 在不同层中，保持patch数量和patch大小（16*16）不变。
• 一整张图就是一个窗口，即VIT是在整张图的维度上计算attention的。

我们再来看图（a），它表示Swin Transformer的架构，具有以下特点：

• 在窗口内计算attention。注意到在每一层中，它将图片切分成若干个窗口（红框），各个窗口内含有固定数量的patch（灰框，实操中，patch_size一般设为4*4，每个window内的patch数量一般设为49，演示图中patch数量为16），各个窗口内独立计算attention（而不是基于全图）。这项优化对应解决上文说的“高分辨率图像引起计算复杂度上升”的问题。
• 在层间做patch合并以获得多尺度特征。这个操作在论文中被称为patch merging。在这一步中，我们会通过某种方法，改变原始图片大小，然后继续按相同的方式划分窗口（此时每个窗口内依然有49个patch，每个patch对应的原始图片的尺寸还是4*4，只是窗口的数量变少了）。通过这种方式，我们模拟了CNN架构中产生不同尺寸特征图的过程，以此更好解决CV中的预测密集型任务（检测、分割等）。这项优化对应解决上文说的“同一实体尺寸变化”的问题。

看到这里如果觉得迷惑，也没关系，我们会在后文对这几块做详细阐述。不难发现，Swin Transformer做的这两项优化，是有针对性的：都是针对CV任务做的专门优化。如果你再仔细感受一下，可能会发现它甚至借鉴了CNN架构对图像任务的归纳偏置（一种先验假设）的契合方法（参见这篇文章6.2.1） 。而在NLP任务中，其实并不关心多尺度特征这类问题。由于Swin Transformer这种对图像任务的针对性，导致它无法像VIT那样用统一的方式表示不同模态的特征，这也是文章作者留给后人填的一个坑。

好，讲完了Swin Transformer在图像任务上做的针对性改进，接下来我们整体看一下Swin Transformer的架构和运作流程。在掌握整体流程之后，我们再分块对细节进行详细阐述。

二、Swin Transformer的整体架构和运作流程

我们先来看（a），它描绘了Swin Transformer的整体架构和前向过程。

（1）首先，假设我们有一张224*224*3的输入图片，即这里H=224, W=224

（2）然后，划分patch，即图中所示的patch partition。实操中我们设置patch的大小为4*4，如果把一张图想象成一个长方体。那么划分后，将得到56*56个大小为4*4*3的小方块，如果我们把这些小方块拉平，就得到56*56*48的立方体，也就是图中在patch partition后所示的

但是，如果我们想让通道数变大点，比如从48变到96，这时我们就可以再增加一个linear embedding层来做线性转换。也就是图中stage1的linear embedding过后，最终的图片大小为56*56*96。

在实际代码操作中，以上两个步骤可以用一个nn.Conv2d完成，我们在VIT原理篇中画过示例，这里我们再展示一次，将图中小patch的尺寸换成4*4*3，将卷积核变成96个4*4*3的卷积核，stride=patch_size=4，padding=0，就能得到一个56*56*96大小的方块了。那么拉平后，类比于NLP任务，我们就可得到56*56=3136个token，每个token的维度为96。

（3）和VIT一样，秉持着一个patch就是一个token的原则，我们可以把数据输入到Swin Transformer中了。Swin Transformer具体的样子在（b）中已绘制出来，我们放在后文细谈。经过Swin Transformer block的处理，我们得到输出结果，其尺寸为56*56*96，对应着演示图中stage1的输出结果

（4）好，现在我们进入Stage2了，这时我们要做一个重要的操作：Patch Merging。我们在第一部分中已经说过它的大致作用。经过Patch Merging后，我们的输入数据尺寸变成28*28*192，再经过Swin Transformer block后的输出结果尺寸也一样，对应着图中。关于patch merging的细节，我们也在下文介绍。

（5）以此类推，stage3的输入输出尺寸为14*14*384，stage4的输入输出尺寸为7*7*768。拿着stage4的输出特征，我们就可以做后续的处理，用到下游的分割、检测、分类任务上了。

现在，我们回头再来看（b）

图（b）描绘了Swin Transformer block做的事情：输入数据先正常在窗口范围内做attention（W-MSA, W表示Window） ，然后经过LN和MLP一系列模块后，又做了一次移动窗口的attention（SW-MSA， SW表示Shifted Window） ，然后才是最终的输出结果。因此你可以发现演示图中的block都是偶数的。因为两个相连的block要依次完成W-MSA和SW-MSA的步骤。这个SW-MSA，就是Swin Transformer的灵魂所在了（连名字都是这么来的），大白话就是移动窗口下的attention计算。我们在后文来详细看它的实现细节。

好，目前为止，我们把Swin Transformer的整体流程过完了，我们回顾一下其中涉及到的几个细节问题：

• patch分割（patch partition）怎么做？
• patch融合（patch merging）怎么做？
• Swin Transformer Block中的W-MSA和SW-MSA怎么做？

其中，第一个问题我们在本部分中已给出了详细的解答。所以我们集中在后面两个问题上。

三、Patch Merging

3.1 Patch Merging操作方法

配合图例，我们来看patch merging的详细过程。我们以第二部分流程图中stage1的输出结果为例。

上图中HWC的图片对应到stage1的输出结果，尺寸就是56*56*96，为了绘图方便，C维度没有画出。

不难发现此时每一个小格子即表示一个patch，我们对patch进行编号，接下来我们就可以按照设定好的编号，把固定位置的patch取出来，拼成四份小图。

接着，我们在C维度concat这四份小图，生成一份尺寸为H * W * 4C的数据。以上两步骤的代码操作如下：

        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)

最后，我们引入线性层 nn.linear(4C, 2C) ，减少通道数，使得最终输出的数据维度为H/2 * W/2 * 2C。

完整的代码流程如下，大家注意看注释，注意看数据的尺寸变化：

class PatchMerging(nn.Module):
    r""" Patch Merging Layer.

    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def forward(self, x):
        """
        x: B, H*W, C
        """
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"
        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."

        x = x.view(B, H, W, C)

        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C

        x = self.norm(x)
        x = self.reduction(x)

        return x

    def extra_repr(self) -> str:
        return f"input_resolution={self.input_resolution}, dim={self.dim}"

    def flops(self):
        H, W = self.input_resolution
        flops = H * W * self.dim
        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
        return flops

3.2 Patch Merging注意事项

关于patch merging，这里再强调两点。

(1）做merge的本质上是patch，而不是一个pixel向量。

你可能想问，在3.1的图中，为什么那些红黄蓝绿的小方框，表示的是一个patch，而不是图中的一个像素点向量呢？

我们回到第二部分对patch partition的介绍图中，可以发现，partition后图中每个位置上的向量（维度为1*1*96），其实是由原始图片中一块4*4*3的patch经过Conv2d处理而来。

所以，这些小方框，在处理过后的图中可理解成一个1*1*96维度的pixel向量，但是本质上，它对应着原图中的一个4*4*3的patch。 经过patch merging的concat和linear操作后，它可能不严格对应着原图中的某块patch，而是若干块patch在通道维度拼接而成，但抽象意义上，它仍表示4*4的patch范围。

那么经过patch merging后得到的 H/2 * W/2 * 2C 的特征图，还要再进行patch partition吗？不需要了，这一点在第二部分整体架构流程图中也有体现。因为正如上面所说，此时特征图中每个1 * 1 * 2C的向量，其抽象表示范围依然对应着最原始输入图片中4*4的patch。所以这里，我们只需正常再做window划分（下文会说），而不用再去做patch partition了。

所以我们又说，在Swin Transformer中，patch是基本计算单位，这一点不仅表现在它类比于token参与Transformer block计算，也表现在它的抽象意义上。现在回头看1.2中的（a）图，是不是理解更深刻了些？

（2）patch merging和CNN中的池化操作相似

不难发现，其实patch merging就非常像CNN中的池化操作，只是这里将池化操作常用的min/max/average替换成了按固定位置取patch。

四、W-MSA与SW-MSA

好，接下来，我们就来看Swin Transformer是如何在窗口范围内做attention的。

再回顾一下标题中的两个英文符号的含义：

• W-MSA：Window-Multi Head Attention，表示在窗口范围内做attention
• SW-MSA：Shifted Window-Multi Head Attention，表示在移动窗口的范围内做attention

我们先来看W-MSA和SW-MSA是什么。再来看相比于VIT的全局Attention，这两个操作带来的好处是什么。

4.1 W-MSA

4.1.1 W-MSA的流程

W-MSA的流程如左图所示。在VIT中，左图中灰色的patch类比于一个个token，我们对全部的token做attention计算。在Swin Transformer中，我们先划分窗口（红框），每个窗口内默认有49个patch。我们只在窗口的范围内做attention

4.1.2 W-MSA的缺点

但是在窗口内计算attention会有一个明显的缺点：patch的感受野变小了。

什么意思呢？就是原来在VIT里，一个patch可以和全局的patch做attention，它能看到整张图的情况。而现在却只能看到窗口内的视野了。为了解决这个缺陷，我们引入了SW-MSA，即在移动窗口的范围内做attention。

4.2 朴素SW-MSA

4.1.1的右图刻画了一种朴素粗暴的移动窗口的办法，你可以理解成将原始的窗口向右下角移动2个patch，就可得到图中的右图。

这样移动的好处是，对一些patch来说，两次计算attention时，它们分属不同的窗口，这样可以让它们和不同的patch交互，扩大感受野。

但这样做的缺点也是明显的：原来我就4个窗口，每个窗口内固定都是49个patch。经过你这么移动，我有9个窗口，每个窗口内的patch数量还不一致，这不是加重了模型的计算复杂度，以及我写代码的难度吗？

所以，作者在这里创新性地提出了一种环形移动窗口法（cyclic shift） 。

4.3 环状SW-MSA

这个“环状”听起来玄学，其实理解起来不复杂。如上图，最开始是9个不均匀的patch。接着，为了让它均匀，我们将原来的A、B、C部分做一个旋转拼接，就又可以恢复成一个有4个window，且patch数量均匀的数据了。

那么恢复过后呢？还是直接在窗口内计算attention吗？当然不是，这样做会有一个显著的问题：假设这张图片绘制的是蓝天和大地，其中C中包含了蓝天的信息。那么此时在左下角那个窗口中，如果你做attention，势必就会让大地和天空一起做attention，那么这时模型大概率可能会学到天空在大地的下面，这当然是不合理的。

那怎么解决这个问题呢？作者在这里创新性地提出了masked MSA（掩码MSA） ，也就是上图大括号后面的部分（这个部分上图画得比较粗糙，我们在后文4.4中会给出更详细的说明）

做完masked MSA后，我们再把A、B、C还原回原来的位置，就可以了。

4.4 Masked Attention

以上图例来自Swin Transformer github issue，由一作绘制得出，它非常直观地展现出了掩码attention的计算过程。我们可以将其和4.3中做完环状移动窗口后的图对比来看。

先看Window0，在这个窗口中，patch没有任何拼接处理（左图），因此它不需要做任何掩码操作，可以正常做attention。右图中window0的颜色为全黑，意味着无需任何掩码操作。

再看Window2，在这个窗口中，标号为3的部分来自同一批patch，标号为6的部分来自环状移动过来的另一批patch。因此3和6之间是不能做attention的，这意味着如果attention score是由3的一块patch和6的一块patch计算得来，我们就需要把这个score设置成-100，这样一来在后续做softmax时，对应位置的结果就可以小到忽略不计，以此来取得遮掩（masked）效果。右图中黄window2的黄色部分表示不需要做mask的分数，黑色部分表示需要做mask的分数。

如果你还有疑惑，那么再来看一眼下面这张图，应该可以帮你解答疑惑：

同理我们可以推知Window1和Window3的结果。

4.5 复习 Swin Transformer Block

现在我们再回头复习一遍Swin Transformer Block的运作流程。

（1）首先，来了一排可以被当作token的patch，我们对这些patch划分好了Window。

（2）先做W-MSA，在窗口范围内计算attention

（3）再做SW-MSA，以环状的方式移动窗口，目的是为了让窗口数量和W-MSA保持一致。移动窗口的目的是为了扩大每个patch的感受野。

（4）因为采用了环状移动的方式，导致patch间的相对位置变动了，因此我们引入mask attention，对被打乱位置的patch间的attention score做掩码操作。

（5）因为每次都必须进行W-MSA + SW-MSA的操作，因此block的数量必须为偶数。

好！到此为止，我们就把整个Swin Transformer的核心技术说完了。但别忘了，还有个重要问题没解决：为什么要从全局attention变成窗口attention呢？

五、窗口attention计算量分析

我们直接对上面这个问题抛出回答：相比于全局attention，窗口attention能有效降低计算量。

那接下来我们肯定就要问，那到底能降多少呢？这点在论文中只给出了结论，而没有给出详细的推导过程。接下来，我们就来看看计算推导细节。

5.1 矩阵相乘的计算量

我们一般用FLOPs（floating point operations，浮点运算次数） 来表示运算量的大小。对于“两矩阵相乘”这个操作而言，其运算量 = 乘法运算的次数 + 加法运算的次数。

来看一个具体例子：

两矩阵相乘，为了获取图中深橘色部分的元素，我们一共需要进行n次乘法运算和n-1次加法运算。

那么现在结果矩阵中，一共有m*p个橘色方块，则意味着我们需要进行：m*p*(n + n - 1)次浮点计算。

再进一步，假设此时在蓝色和绿色的矩阵外，我们还有一个bias矩阵，意味着计算单个橘色方块时我们需要进行n次乘法和n-1+1次加法运算，那么此时总计算量为：m*p*(n+n) = 2mnp。当然，即使不加这个bias，我们也可以把-1项给忽略，得到相同的结果。

所以这里我们总结下，假设有两个矩阵A和B，它们的维度分别为(m, n)和(n, p)，则这两矩阵相乘的运算量为2mnp。

一般在矩阵运算中，**乘法运算的时间要高于加法运算的时间，因此有时在统计运算量时，我们只考虑乘法运算的次数，则此时两矩阵相乘的运算量可近似为mnp，而Swin Transformer的论文中采用的就是这种近似。

5.2 全局attention计算量

有了5.1的结论，我们就可以来算attention的计算量了，毕竟attention说白了也是一系列矩阵操作。我们先来看全局attention的计算量。

我们忽略batch维度（认为b = 1），则此时输入数据的维度为(hw, C)，其中hw（即h * w）表示patch数量，也就等同于token序列的长度，C表示token embedding维度。

（1）计算Attention的第一步，我们需要让输入数据和这三个矩阵相乘，得到Q、K、V三个结果矩阵。而这三个矩阵的维度都是(C, C)，那么按照5.1中给出的公式：(hw, C) * (C, C) = (hw, C)，此时的运算量为

(2) 接下来，我们需要将Q和K相乘，得到attention score。则：(hw, C) * (C, hw) = (hw, hw)，此时的运算量为

(3)接下来，我们要将attention score和V相乘，则：(hw, hw) * (hw, C) = (hw, C)，此时的运算量为

（4）最后，我们要将结果过一层线性映射层，则：(hw, C) * (C, C) = (hw, C)，此时运算量为

以上四步的运算量加总，则全局attention的总运算量为

5.3 窗口attention计算量

前面我们说过，一个窗口内固定有49个patch，也就是窗口可以拆成7*7的矩阵。

我们先来算一个窗口内做attention的计算量，然后再乘上窗口总数，就可以得到总计算量了。

我们设窗口的H和W的值为：H=M，W=M，在默认操作中M=7，则代入5.1的公式可求得单个窗口attention的计算量为 。

而此时，我们的窗口总数为： 。

则在窗口下，attention的计算总量为： 

5.4 全局MSA VS 窗口MSA

经过5.2和5.3这么一算，两者的计算量就出来了：

你看，在全局attention中，计算量和图像尺寸大小h*w是平方关系，但是变成窗口计算之后，就变成线性关系了。 这大大降低了运算复杂度，这也使得Swin Transformer能够被用于处理高分辨率的图像。

如果你对复杂度还没有那么直观的感受，你可以将M=7，h=w=56带入上面的式子算一算，就能感受到窗口MSA的效果了。

总结来说，通过“窗口MSA + 移动窗口”这种方式，Swin Transformer既做到了节省计算量，又做到了扩充每个patch的感受野。而通过“patch merging”这种方式，它模仿了CNN架构中输出不同尺寸特征图的过程，使得模型能够好处理预测密集型任务（检测、分割）。

关于实验部分，这里就不展开细说了，感兴趣的朋友可以翻阅原始论文。

六、参考

1、https://arxiv.org/pdf/2103.14030.pdf

2、https://github.com/microsoft/Swin-Transformer

3、https://www.bilibili.com/video/BV13L4y1475U/