Press "Enter" to skip to content

Transformer论文解读五(Swin Transformer)

本文中要介绍的Swin Transformer全名为Swin Transformer: Hierarchical Vision Transformer using Shifted Windows,是2021年微软研究院发表在ICCV上的一篇文章,在多项视觉任务中展现出优异性能。

 

1. 改进

 

相比于Vision Transformer(讲解),Swin Transformer做出了几点改进:

 

 

    1. 层次化构建方法 :使用了类似卷积神经网络中的层次化构建方法(Hierarchical feature maps),比如特征图尺寸中有对图像下采样4倍的,8倍的以及16倍的,这样的backbone有助于在此基础上构建目标检测,实例分割等任务。而Vision Transformer中直接下采样16倍,后面的特征图也是维持这个下采样率不变。

 

    1. 特征图划分 :在Swin Transformer中使用了Windows Multi-Head Self-Attention(W-MSA)的概念,比如在下图的4倍下采样和8倍下采样中,将特征图划分成了多个不相交的窗口,并且Multi-Head Self-Attention只在每个窗口内进行。相对于Vision Transformer中直接对整个特征图进行Multi-Head Self-Attention,这样做的目的是 减少计算量 ,尤其是在浅层特征图很大的时候。

 

    1. 窗口变换 :特征图划分虽然减少了计算量,但会隔绝不同窗口之间的信息传递,因此所以在论文中作者又提出了Shifted Windows Multi-Head Self-Attention(SW-MSA)的概念,目的在于使 信息在相邻窗口中传递 。

 

 

 

2. 网络架构

 

论文中给出了Swin Transformer(Swin-T)网络的架构图。可以看出整个框架的基本流程如下:

 

 

 

    1. 图片分块 :首先将图片输入到Patch Partition模块中进行分块,即每 4 × 4 4\times4 4 相邻的像素为一个Patch,然后在channel方向展平。由于每个patch就有 4 × 4 = 16 4\times4=16 个像素,每个像素有R、G、B三个值,所以展平后的深度是 16 × 3 = 48 16\times3=48 1 6 × 4 8 ,所以通过Patch Partition后图像shape由 [ H , W , 3 ] [H, W, 3] [ H , W , 3 ] 变成了 [ H 4 , W 4 , 48 ] [\frac{H}{4}, \frac{W}{4}, 48] [ , , 4 8 ] 。

 

    1. 线性变换 :通过Linear Embeding层对每个像素的channel数据做线性变换,由 48 48 4 8 变成 C C C ,即图像shape再由 [ H 4 , W 4 , 48 ] [\frac{H}{4}, \frac{W}{4}, 48] [ , , 4 8 ] 变成了 [ H 4 , W 4 , C ] [\frac{H}{4}, \frac{W}{4}, C] [ , , C ] 。在源码中通过一个卷积层实现。

 

    1. 网络模块堆叠 :通过四个Stage构建不同大小的特征图,除了Stage1中先通过一个Linear Embeding层外,剩下三个stage都是先通过一个Patch Merging层进行下采样。然后都是重复堆叠Swin Transformer Block。注意,Transformer Block有两种结构,如图(b)中所示,这两种结构的不同之处仅在于一个使用了W-MSA结构,一个使用了SW-MSA结构。而且这两个结构是成对使用的,先使用一个W-MSA结构再使用一个SW-MSA结构。所以堆叠Swin Transformer Block的次数都是偶数(因为成对使用)。

 

    1. 分类头 :最后,对于分类网络,在后面接上一个Layer Norm层、全局池化层以及全连接层,得到最终输出。

 

 

3. Patch Merging

 

除了Stage1以外的Stage中,首先要通过一个Patch Merging层进行下采样,如下图所示

 

假设输入Patch Merging的是一个4×4大小的单通道特征图,Patch Merging的过程如下:

 

 

    1. 将每个2×2的相邻像素划分为一个patch

 

    1. 将每个patch中相同位置(同一颜色)像素拼在一起,就得到了4个feature map

 

    1. 将四个feature map在深度方向进行concat拼接

 

    1. 通过一个LayerNorm层

 

    1. 通过一个全连接层在feature map的深度方向做线性变化,将feature map的深度由C变成C/2。

 

 

可以看出,通过Patch Merging层后,feature map的高和宽会减半,深度会翻倍。

 

4. W-MSA

 

引入Windows Multi-head Self-Attention(W-MSA)模块是为了减少计算量。

 

如下图所示,在Self-Attention计算过程中,普通的Multi-head Self-Attention(MSA)模块将feature map中的每个patch和所有的像素去计算。但Windows Multi-head Self-Attention(W-MSA)模块首先将feature map按照MxM(例子中的M=2)大小划分成一个个Windows,然后单独对每个Windows内部进行Self-Attention。

 

 

4. SW-MSA详解

 

采用W-MSA模块时,只会在每个窗口内进行自注意力计算,所以窗口与窗口之间是无法进行信息传递的。为了解决这个问题,作者引入了Shifted Windows Multi-Head Self-Attention(SW-MSA)模块,即进行偏移的W-MSA。

 

如下图所示,W-MSA和SW-MSA是成对使用,假设第L层使用W-MSA,那幺第L+1层使用的就是SW-MSA。根据左右两幅图对比能够发现窗口(Windows)发生了偏移。偏移后的窗口使得之前的临近窗口之间进行交流,就解决了不同窗口之间无法进行信息交流的问题。

 

但是,可以发现通过将窗口进行偏移后,原来的4个窗口变成9个窗口了。为了解决这个麻烦,作者又提出而了Efficient batch computation for shifted configuration,一种更加高效的计算方法:

 

主要思想是将左上角的小窗口移动到右下角,和右下角的小窗口合并在一起,这样又和原来一样是4个4×4的窗口了,所以能够保证计算量是一样的。

 

为了防止将不同的区域合并在一起导致信息混淆的问题,在实际计算中使用的是masked MSA,通过设置蒙板隔绝其他区域的信息,使得每个像素只与自己所在区域的像素进行计算,从而解决了这个问题。

 

参考文献

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注