【论文精读】Swin Transformer

type

status

password

date

slug

summary

1、概述

1.1、动机

目前将 Transformer 从自然语言处理领域应用到计算机视觉领域主要有两大挑战：

视觉实体的方差较大，例如同一个物体，拍摄角度不同，转化为二进制后的图片就会具有很大的差异。又如不同尺寸下，人像照片有很大的不同（1寸免冠1、1080p的风景照）。

图像分辨率高，像素点多，如果采用ViT模型，自注意力的计算量会与像素的平方成正比。

1.2、解决方案

针对上述两个问题，论文中提出了一种基于滑动窗口机制，具有层级设计（下采样层） 的 Swin Transformer。

其中滑窗操作包括不重叠的 local window，和重叠的 cross-window。将注意力计算限制在一个窗口（window size固定）中，一方面能引入 CNN 卷积操作的局部性，另一方面能大幅度节省计算量，它只和窗口数量成线性关系。通过下采样的层级设计，能够逐渐增大感受野，从而使得注意力机制也能够注意到全局的特征。

在论文的最后，作者也通过大量的实验证明Swin Transformer相较于以前的SOTA模型均有提高，尤其是在ADE20K数据和COCO数据集上的表现。也证明了Swin Transformer可以作为一种通用骨干网络被使用。

2、模型结构

整个模型采取层次化的设计，一共包含 4 个 Stage，除第一个 stage 外，每个 stage 都会先通过 Patch Merging 层缩小输入特征图的分辨率，进行下采样操作，像 CNN 一样逐层扩大感受野，以便获取到全局的信息。

2.1、数据流

首先将RBG图像分割为不重叠的patches，然后把patch拉直作为“token”，patch大小为4x4，所以每一个token维度为4x4x3=48，然后做linear embedding，将特征映射到任意维度（设为C），假设patches数量为N (N=H/4 * W/4)，则得到维度为NxC的矩阵。

在每个 Stage 里（除第一个 Stage ），均由Patch Merging和多个Swin Transformer Block组成。

其中Patch Merging模块主要在每个 Stage 一开始降低图片分辨率，进行下采样的操作。
而Swin Transformer Block具体结构如右图所示，主要是LayerNorm，Window Attention ，Shifted Window Attention和MLP组成。

代码

在微软亚洲研究院提供的代码中，是将Patch Merging作为每个 Stage 最后结束的操作，输入先进行Swin Transformer Block操作，再下采样。而最后一个 Stage 不需要进行下采样操作，之间通过后续的全连接层与 target label 计算损失。

💡

ViT 会使用[CLS]作为分类的 token。而 Swin-T 则是直接做平均（avgpool），输出分类，有点类似 CNN 最后的全局平均池化层。

2.2、位置编码

ViT 在输入会给 embedding 进行位置编码。而 Swin-T 是在计算 Attention 的时候做了一个相对位置编码，我认为这是这篇论文设计最巧妙的地方。

2.5.3、相对位置编码

2.3、Patch Embedding

在输入进 Block 前，我们首先通过Patch Partition将原始图片裁成一个个patch （patch大小为window_size * window_size，即窗口大小），然后利用Linear Embedding进行嵌入。

💡

这里可以通过二维卷积层，将 stride，kernel_size 设置为 window_size 大小。设定输出通道来确定嵌入向量的大小。最后将 H,W 维度展开，并移动到第一维度。

论文中输出通道设置为48，但是代码中为96，以下我们均以代码为准。Batch_size=128

代码

2.4、Patch Merging

该模块的作用是在每个 Stage 开始前做降采样，用于缩小分辨率，调整通道数进而形成层次化的设计，同时也能节省一定运算量。

💡

在 CNN 中，则是在每个 Stage 开始前用stride=2的卷积/池化层来降低分辨率。

每次降采样是两倍，因此在行方向和列方向上，间隔 2 选取元素。然后拼接在一起作为一整个张量，最后展开。此时通道维度会变成原先的 4 倍（因为 H,W 各缩小 2 倍），此时再通过一个全连接层再调整通道维度为原来的两倍。

下面是一个示意图（输入张量 N=1, H=W=8, C=1，不包含最后的全连接层调整）

代码

2.5、Swin Transformer Block

2.5.1、Window Partition/Reverse

window partition函数是用于对张量划分窗口，指定窗口大小。将原本的张量从 N*H*W*C, 划分成 num_windows*B, window_size, window_size, C，其中 num_windows = H*W / window_size*window_size，即窗口的个数。而window reverse函数则是对应的逆过程。这两个函数会在后面的Window Attention用到。

2.5.2、W-MSA

传统的 Transformer 都是基于全局来计算注意力的，因此计算复杂度十分高。而 Swin Transformer 则将注意力的计算限制在每个窗口内，进而减少了计算量。我们先简单看下公式

主要区别是在原始计算 Attention 的公式中的 Q,K 时加入了相对位置编码。

2.5.3、相对位置编码

相对位置编码矩阵直接加在了attention 矩阵中，注意的维度和attention矩阵相同，即seq_len x seq_len。

相对位置取决于参考点在哪，以一个2x2大小，共4个像素的图像为例，其相对位置信息有四种

2.5.6、SW-MSA

前面的 Window Attention 是在每个窗口下计算注意力的，为了更好的和其他 window 进行信息交互，Swin Transformer 还引入了滑动窗口操作。

SW-MSA是紧接在W-MSA后面的，在第层（左侧），采用W-MSA，并在每个窗口内计算自注意力，如左2图所示。在接下来的第层（右侧），采用SW-MSA，窗口向右下角移位3格，从而形成了新的窗口，如左3、4图所示。移位后的窗口包含了原本相邻窗口的元素，达到了窗口之间信息交互的作用。

在代码实现中通过torch.roll实现

💡

但这也引入了一个新问题，即红色、绿色部分如何计算（这两部分占窗口的一半）。在实际代码里，我们是通过对特征图移位，并给 Attention 设置 mask 来间接实现的。能在保持原有的 window 个数下，最后的计算结果等价。