Swin Transformer:用Transformer实现CNN多尺度操作

文本是关于Swin Transformer基础知识的了解

论文：https://arxiv.org/pdf/2103.14030

项目：https://github. com/microsoft/Swin-Transformer.

实现一个Swin Transformer：Swin Transformer模型具体代码实现-CSDN博客

Swin Transformer mlp版本实现：Swin Transformer 变体1：使用MLP代替多头注意力机制-CSDN博客

Swin Transformer v2版本实现：Swin Transformer变体2：Swin Transformerv2模型-CSDN博客

原理总览：

1、总述：

Swin Transformer和Vision Transformer很像，他们都是将NLP领域的Transformer架构引入计算机视觉领域期待大力出奇迹的产物，不同的是，Vision Transformer由于它固定的尺度（16x16大小的patch），它并不适合密集预测，比较适合于图像分类任务，对于目标检测、语义分割这种任务就有点困难，而且最主要的问题是，ViT使用的是全局注意力机制，这种注意力的最大问题是计算成本和内存消耗太大。而Swin transformer的目标不是做一类任务，它想做视觉任务的骨干模型。

Swin Transformer 做出的第一大改进就是在窗口内计算注意力，这样的话就把复杂度从与图片大小的平方关系改进到了线性关系，因为窗口内计算注意力的计算量相比全局注意力的计算量差距很大，很多个窗口完全可以并行在batch中计算。

第二大改进就是引入了移动窗口，因为Swin Transformer 使用窗口注意力，那么就有一个问题，不同窗口之间的信息怎么交互呢，作者引入了窗口移动机制，使用两个block层，第一个block层中使用现有的窗口做自注意力计算，之后移动窗口，具体来说就是将patch矩阵沿着右下角移动window_size//2个单位，之后在第二个block中计算新的窗口中的注意力计算，在计算完成之后进行反向移动操作，将之前的移位操作复原。

第三大改进就是学习CNN的pooling机制，实现了层次化的多尺度建模，这样的话有助于学习不同尺度的特征，对密集预测十分有帮助，这点在论文的实验结果部分表现得很清楚。

Swin Transformer的这些特性使其与广泛的视觉任务兼容，包括图像分类（在ImageNet-1K上达到87.3%的top-1准确率）和密集预测任务，如目标检测（在COCO测试集上达到58.7的box AP和51.1的mask AP）和语义分割（在ADE20K验证集上达到53.5的mIoU）。其性能超越了之前的最先进水平，COCO上提高了+2.7 box AP和+2.6 mask AP，ADE20K上提高了+3.2 mIoU，展示了基于Transformer的模型作为视觉骨干网络的潜力。层次化设计和移动窗口方法也证明了对全MLP架构的益处。

2、总体架构：

Swin Transformer总体来说分为4个阶段：

首先，对于输入的图像数据，需要先进行Patch Partition操作，也就是分块，这里的patch大小为4x4，相比于ViT的patch大小16x16，Swin Transformer的每个patch可以保留更细节的信息。如上图，(H//4)*(W//4)表示的是patch数量，48表示的是3x4x4，也就是一个patch的三个通道的所有像素的拼接，以此作为每个patch的特征，之后经过嵌入，将48映射为C（tiny和small模型中默认为96，big和large模型默认为192）,这里所有的操作其实在代码实现中都是通过一个卷积层实现的，具体可以参考ViT的实现。

2.1、stage 1：

上图架构是Swin Transformer-T模型的架构图，他的第一阶段有2个block层，在Swin Transformer中，2个block作为一组同时出现，所以每个阶段的block数量都为偶数：

在第一个阶段中，对于输入的数据【B，H//4*W//4，C】：

在第一个block中不进行窗口移动，只进行窗口自注意力计算，我们可以详细分析一下这中间的特征形状变化：对于输入数据形状：【B，H//4*W//4，C】，首先将一维序列展开为二位序列：【B，H//4，W//4，C】，之后划分窗口：【B*num_windows，windows_size，windows_size，C】，其中windows_size是窗口大小，默认为7，num_windows为窗口数量=（H//4//windows_siz，W//4//windows_size），之后根据每个block的注意力头的数量重塑形状为：【B*num_windows，head，windows_size*windows_size，C//head】，其中，windows_size*windows_size就是一个窗口内的序列长度，C//head就是每个元素的特征，这样的话，就把原本ViT的全局注意力计算变为了窗口注意力计算（ViT：【B，head，（H//16）*（W//16）+1，C//head】），这里需要注意，在计算完计算力得分之后，还有一个相对位置偏移需要相加，这一点将在本章的第5节讲解。计算完注意力，需要将窗口合并，即形状变为【B，H//4，W//4，C】，之后将数据形状重塑回到原始形状【B，H//4*W//4，C】即可进入下一个block。

在第二个block中，需要进行窗口的移动，这个操作在划分窗口前进行，只需要调用torch.roll函数即可。在注意力计算阶段，由于窗口发生了变化，相对应的注意力计算区域也发生了变化，为了保证窗口数量不变以及不计算不应该计算位置的注意力得分，这里对于 $QK^{T}$ 的结果要进行掩码处理，具体操作在本章的第6节讲解。进行完注意力计算之后，需要进行窗口合并，即形状变为【B，H//4，W//4，C】，然后将之前移动的窗口还原，之后将数据形状重塑回到原始形状【B，H//4*W//4，C】。

2.2、stage 2：

stage 2第一个操作是降采样操作，也就是patch_merging，这个将在本章的第三节进行讲解，降采样的结果就是将图像尺寸减半，将通道加倍，这个操作是在模仿CNN中的池化操作，经过该操作输入数据变为【B，H//8*W//8，2C】。

stage 2 的Swin Transformer block的流程和stage1 完全一致，就是输入数据形状不一样。

2.3、stage 3：

stage 3第一个操作和stage 2一样，都是降采样把图像尺寸减半，将通道加倍，这个操作是在模仿CNN中的池化操作，经过该操作输入数据变为【B，H//16*W//16，4C】。

stage 3 的Swin Transformer block的流程和stage1 完全一致，就是输入数据形状不一样，同时stage 3有6个block。

2.4、stage 4：

stage 4第一个操作和stage 2一样，都是降采样把图像尺寸减半，将通道加倍，这个操作是在模仿CNN中的池化操作，经过该操作输入数据变为【B，H//32*W//32，8C】。

stage 4 的Swin Transformer block的流程和stage1 完全一致，就是输入数据形状不一样。

2.5、输出处理：

stage 4的输出经过平均池化和压缩之后，得到输出结果，形状为：[B, num_features]，之后经过Swin Transformer head的映射，输出结果[B, num_classes]。

3、patch_merging：

该模块是模型的降采样模块，是把图像尺寸减半，将通道加倍，这个操作是在模仿CNN中的池化操作。下面的图可以简要描述这种操作：

上图中的1，2，3，4代表的是应该在分割后归属的图的标号而不是真实值，按照上图所示的逻辑，将原始数据分割为4部分，每一部分形状为【B，H//2，W//2，C】，在C通道维度拼接为【B，H//2，W//2，4C】，之后将4C维度映射为2C：【B，H//2，W//2，2C】。

4、窗口自注意力机制和全局自注意力机制运算量：

$\Omega (MSA)=4hwC^{2}+2(hw)^{2}C$

$\Omega (WMSA)=4hwC^{2}+2M^{2}hwC$

其中，h和w是图像的高和宽，C是通道数量，M是窗口大小。

我们知道全局多头自注意力机制的运算包括4个矩阵映射、2个矩阵乘法，其中，矩阵映射是K、Q、V（h*w*C）与C*C的矩阵的相乘，运算量为 $3hwC^{2}$ ，当计算完注意力分数之后，要将attn经过映射输出，这部分也是h*w*C与C*C的矩阵的相乘，运算量为 $hwC^{2}$ 。两个矩阵乘法是 $QK^{T}$ 和 $AV$ ，他们都是h*w*C与C*h*w矩阵乘法，运算量都为 $(hw)^{2}C$ ，所以全局注意力运算量总和为 $4hwC^{2}+2(hw)^{2}C$ 。

窗口多头自注意力机制的运算也包括4个矩阵映射、2个矩阵乘法，其中，矩阵映射是K、Q、V（M*M*C）与C*C的矩阵的相乘，运算量为 $3M^{2}C^{2}$ ，当计算完注意力分数之后，要将attn经过映射输出，这部分也是M*M*C与C*C的矩阵的相乘，运算量为 $M^{2}C^{2}$ 。两个矩阵乘法是 $QK^{T}$ 和 $AV$ ，他们都是M*M*C与C*M*M矩阵乘法，运算量都为 $M^{4}C$ ，所以一个窗口注意力运算量总和为 $4M^{2}C^{2}+2M^{4}C$ ,总的窗口个数是（h//M）*（w//M），相乘得到窗口注意力运算量总和为 $4{hw}C^{2}+2M^{2}hwC$ 。

5、相对位置偏移：

为什么要采用相对位置？ViT的结果表明使用相对位置和绝对位置对最后的分类结果影响不大，但是Swin Transformer它的用途不局限于图像分类，在目标检测、语义分割等密集预测任务中，使用相对位置比使用绝对位置的效果要好很多，这一点作者做了消融实验，结果如下：

Swin Transformer的相对位置偏移表大小为【（2M-1）*（2M-1），head_num】，其中，2M-1是窗口的高和宽的相对位置范围，即如果M=3，那么（-2，-1，0，1，2）就是相对位置范围。相对位置偏移表是一个可学习的参数表。

然后使用窗口的高和宽堆叠形成二维网格，之后压缩形成【2，M*M】的网格坐标，之后计算第一个M*M坐标和第二个M*M坐标每一对之间的相对坐标差，得到【2，M*M，M*M】，之后重塑形状为【M*M，M*M，2】并给每一个元素加上M-1使得所有元素均不为负，这是因为这些坐标差之后要作为index在相对位置偏移表中取值，之后将最后一维求和得到【M*M，M*M】，将其作为相对位置索引表。

当计算完 $QK^{T}$ 之后，将相对位置索引表中所有元素对应于相对位置偏移表中的head_num维向量取出并重塑形状为【head_num，M*M，M*M】。

我们可以回忆一下 $QK^{T}$ 的形状：【B，head_num，patch_num，patch_num】，其中patch_num=M*M，那么我们给【head_num，M*M，M*M】扩充第零维之后利用广播就可以实现相对位置偏移和 $QK^{T}$ 的相加。

6、窗口移动：

我们知道Swin Transformer的特殊点就在于其可以移动的窗口，但是移动窗口就会出现一个问题：

一个原始特征图如果有4个窗口，沿着右下角移动window_size//2个单位后会发现出现了9个窗口，我们知道如果窗口高和宽不一致，我们无法使用注意力机制，这时最简单的办法就是填充，将高和宽不一致的窗口用0填充至高和宽一致，这样做很简单，但是相应的计算量就上升了，这对于我们来说是不利的，所以我们换一种思路，引入掩码操作：

上图中，每种颜色覆盖范围代表一个逻辑上形式上的窗口（4个），每种数字标识一个实际的窗口（9个），在同一个颜色的逻辑窗口内可能会存在不同的实际窗口，我们可以给一个逻辑窗口内的所有实际窗口做差值，使得掩码中同一个实际窗口中的元素为0，非该实际窗口中的元素为-100，计算注意力时使用的是逻辑上的窗口，这样能保证窗口数量不变。把这个掩码值加入到 $QK^{T}$ 结果中，做softmax操作就可以不需要计算注意力的地方置为0。