论文阅读笔记：Swin Transformer

前言

论文原文：Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
源码地址：https://github.com/microsoft/Swin-Transformer

本篇论文作者试图扩展Transformer的适用性，使其让NLP可以和CNNs在视觉中一样，作为计算机视觉的通用backbone。但是将Transformer在语言领域的高性能转移到视觉领域的重大挑战主要体现在两种模式的两个差异上：

1.规模：与word tokens不同，视觉元素在规模上可能有很大差异，这是一个在目标检测等任务中受到关注的问题。在现有的基于Transformer的模型中，token都是固定比例的，这一属性不适合视觉应用。
2.计算量：与文本段落中的单词相比，图像中的像素分辨率要高得多。存在许多视觉任务，例如需要在像素级进行密集预测的语义分割，而这对于高分辨率图像上的Transformer来说是很困难的，因为自注意力的计算复杂性是图像大小的二次方。

为了克服这些问题，作者提出了一种通用的backbone，称为Swin Transformer，它可以构造分层特征映射，并且计算复杂度与图像大小成线性关系。如图所示：

可以看出Swin Transformer更贴近于传统的具有多尺寸特征图的CNNs backbone，对图像分别下采样4倍，8倍以及16倍。而Vision Transformer（ViT）一直下采样16倍。多尺寸的特征图有利于不同尺寸大小目标的分割和检测任务（解决了差异1带来的问题）。
Swin Transformer的另一个创新点为将连续的self-attention layers划分成多个窗口，每个窗口单独进行muti-head self-attetion（后文中均简称为MSA）计算，文中称为windows muti-head self-attetion（后文中均简称为W-MSA）计算。在W-MSA后会通过shift window的方式将窗口滑动，使得不同window之间可以信息交互，文中成为shifted windows muti-head self-attetion（后文中均简称为SW-MSA）计算，如下图。通过W-MSA和SW-MSA，成功的降低了计算量并且也保证了各个window之间的信息交互，保证了全局视野（解决了差异2带来的问题）。

那么上述的两种方法具体是如何实现的呢？下文会通过介绍Swin Transformer的网络结构来详细讲解。

网络结构

详细的网络结构参数如下表：

上表表示有Swin-T、 Swin-S、 Swin-B、 Swin-L四种尺寸的网络结构，C代表stage1之后的输出通道：

论文中给出的是Swin Transformer的Swin-T的网络结构，如图(a)，其中的Swin Transformer Block如图(b),即将VIT attention Encoding Block中的MSA换成了W-MSA和SW-MSA：

首先将 $H\times W\times 3$ 的图片输入到Patch Partition中进行分块。每个Patch的大小为 $4\times 4\times 3=48$，通过Patch Partition后 shape 从 $H\times W\times 3$ 变为 $\frac{H}{4}\times \frac{W}{4}\times 48$ ，然后在通过Linear Embeding层对每个像素的channel数据做线性变换，变为 $\frac{H}{4}\times \frac{W}{4}\times C$。接着通过4个Stage进行下采样，除了Stage1是Linear Embedding加一对Swin Transformer Block，其他三个Stage都是一个patch merging加若干对Swin Transformer Block。Patch Partition加Stage1的Linear Embedding，这个过程类似于VIT中的Linear Projection of Flattened Patches操作，即patch embedding过程（可查看博文【论文阅读笔记：Vision Transformer】和博文【Vision Transformer（Pytorch版）代码阅读注释】了解）。一对Swin Transformer Block如图(b)中的包含W-MSA的Block和SW-MSA Block，所以Stage中的Block都是2的倍数。

网络细节

Patch Merging

Patch Merging类似于YoloV5中的Focus模块（Focus模块的介绍可查看博文【从YOLOv5源码yolo.py详细介绍Yolov5的网络结构】），只不过在Patch Merging模块在Focus模块之后再进行LayerNorm和通道上的全连接，使得$H=\frac{H_0}{2},W=\frac{W_0}{2},C=C_0\times 2$。如下图：

W-MSA

引入W-MSA的目的是为了减少计算量，但同时也会使得window之间无法进行信息交互。

论文中提到一个 $h\times w\times C$ 的 特征图，MSA 的计算量为公式（1），拆分成 window 宽高均为 M 以后的 W-MSA计算量为公式（2）：

这两个计算量公式是根据Muti-Head Self-Attention公式得来的（Attention公式介绍可查看博文【论文阅读笔记：Attention Is All You Need】）：

两个矩阵相乘（$A^{a\times b}\times B^{b\times c}=C^{a\times c}$）的计算量为：
$$FLOPs=a\times b\times c+a\times (b-1)\times c\approx 2\times a\times b\times c$$
其中包含 $a\times b\times c$ 次乘法和 $a\times (b-1)\times c$ 次加法。

MSA计算步骤如下:
1.由于 $h\times w\times C$ 的 特征图相当于有 $h\times w$ 个 $C$ 维的token向量，将其表示为矩阵：$A^{hw\times C}$

2.Token 矩阵通过乘上 $W_q^{C\times d_k}，W_k^{C\times d_k}，W_v^{C\times d_v}$ 获得对应的 $Q^{hw\times d_k}，K^{hw\times d_k}，V^{hw\times d_v}$，即：

依据Vision transformer的源码，假设 $d_k=d_v=\frac{C}{head}$，head 为 MutiHead 中的 head 个数，因此 3 对矩阵相乘计算量为 $\frac{6hw{C}^2}{head}$。

3.接着计算$Q^{hw\times \frac{C}{head}}\times (K^T{)}^{\frac{C}{head}\times hw}$，得到 $hw\times hw$ 大小的矩阵，计算量为$\frac{2(hw{)}^{2}C}{head}$。

4.除以$\sqrt{d_k}$再计算softmax，论文中提出忽略这部分的计算量：

矩阵大小依然为 $hw\times hw$ 。

5.将得到的 $hw\times hw$ 的矩阵再乘上$V^{hw\times \frac{C}{head}}$，得到 $hw\times \frac{C}{head}$ 大小的矩阵。计算量为$\frac{2(hw{)}^{2}C}{head}$。

6.最后再乘上融合矩阵$W_O^{\frac{C}{head}\ast C}$将特征矩阵还原成 $hw\times C$ 大小的矩阵，计算量为$\frac{2hw{C}^2}{head}$。

总计算量为$\frac{8hw{C}^2}{head}+\frac{4(hw{)}^{2}C}{head}$，因为 muti-head self-attention 中 head ≥ 2,所以总计算量 ≤ $4hw{C}^2+2(hw{)}^{2}C$，这里取最大值进行比较。

W-MSA计算步骤如下:
1.将 $h\times w\times C$ 的 特征图划分到 $\frac{h}{M}\times \frac{w}{M}$ 个宽高均为 M 的 windows 中。

2.将每个宽高为 M 的 windows 进行WSA计算，每个 windows 的计算量为 $4(MC{)}^2+2(M{)}^{4}C$。

3.$\frac{h}{M}\times \frac{w}{M}$ 个 windows 的总计算量为 $\frac{h}{M}\times \frac{w}{M}\times (4(MC{)}^2+2(M{)}^{4}C)=4hw{C}^2+2M^{2}hwC$.

由于W-MSA只和 $hw$ 的一次方成线性关系，而MSA会包含 $hw$ 的二次关系，因此W-MSA大大的降低了计算量。

SW-MSA

前文介绍W-MSA在划分窗口时会带来windows间信息无法交互的问题，所以作者提到使用 shifted window 的方法来增加信息交互，论文的描述如下：

即先从左上角开始使用常规的W-MSA进行窗口划分，每个窗口大小为 $M\times M$，接着以 $stride=(\lfloor \frac{M}{2}\rfloor ,\lfloor \frac{M}{2}\rfloor )$ 滑动窗口。

以上图为例，左边为 $Layer$ $l$ 层 feature map 大小为 $8\times 8$，window 大小为 $4\times 4$，共 $2\times 2$ 个windows。每个window以$stride=(\lfloor \frac{M}{2}\rfloor ,\lfloor \frac{M}{2}\rfloor )=(2,2)$ 滑动，得到右边 $Layer$ $l+1$ 层的 feature map。其过程可以用下图表示：

Efficient batch computation for shifted configuration

SW-MSA 将原本信息不交互的$\lceil \frac{h}{M}\rceil \times \lceil \frac{w}{M}\rceil$ 个 windows 做了信息交互并变成了$(\lceil \frac{h}{M}\rceil +1)\times (\lceil \frac{w}{M}\rceil +1)$ 个，如果对每个 windows 都做 MSA 计算，那么计算量又会比 W-MSA 多，而且 每个 window 的大小也不一样， 无法并行计算。为了解决这个问题，论文提出了一种相邻非重叠窗口之间的连接方式j，如下图：

该过程如果不理解可以看下图：

其将A、B、C三个框中的 window 移动到四个 $4\times 4$ 红色框的对应位置，使其凑成四个 $4\times 4$ 的window。由于有几个 window 是由不相邻的子窗口组成，需要通过Masked MSA 掩膜计算来限制每个 window 中的不同子窗口的 MSA。

至于 window 是如何移动以及掩膜计算如何实现，请看博文：Swin Transformer代码阅读注释。

Relative Position Bias

论文使用了相对位置偏置，公式如下：

论文中只是提到使用改偏置的效果，并没有产出细节，细节请看博文：Swin Transformer代码阅读注释。
使用了相对位置偏置$(rel.pos.)$以后带来了明显的提升。

源码解析

请看博文：Swin Transformer代码阅读注释