计算机视觉Swin-Transformer 网络结构解析

Posted 2023-04-06 笃℃

【计算机视觉】Swin-Transformer 网络结构解析

文章目录

• 【计算机视觉】Swin-Transformer 网络结构解析
• • 1. 介绍
  • 2. 整体框架
  • 3. 模型详解
  • • 3.1 Patch Merging
    • 3.2 Windows Multi-head Self-Attention（W-MSA）
    • • 3.2.1 MSA（Multi-head Self-Attention） 计算量
      • 3.2.2 W-MSA（Windows Multi-head Self-Attention）计算量
    • 3.3 Shifted Windows Multi-Head Self-Attention（ SW-MSA）
    • 3.4 Relative Position Bias
  • 4. 模型结构参数
  • 参考

1. 介绍

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

Swin Transformer是2021年微软研究院发表在ICCV上的一篇文章，并且已经获得ICCV 2021 best paper的荣誉称号。

• Swin Transformer网络是Transformer模型在视觉领域的又一次碰撞。
• 该论文一经发表就已在多项视觉任务中霸榜。

2. 整体框架

在模型详解之前，先来简单对比下 Swin Transformer 和之前的 Vision Transformer（如果不了解Vision Transformer的，可以参考：Vision Transformer）。下图是Swin Transformer文章中给出的图1，左边是本文要讲的Swin Transformer，右边是之前讲的Vision Transformer。通过对比至少可以看出两点不同：

• Swin Transformer 使用了类似卷积神经网络中的层次化构建方法（Hierarchical feature maps），
  • 比如特征图尺寸中有对图像下采样4倍的，8倍的以及16倍的，这样的backbone有助于在此基础上构建目标检测，实例分割等任务。
  • 而在之前的Vision Transformer中是一开始就直接下采样16倍，后面的特征图也是维持这个下采样率不变。
• 在Swin Transformer中使用了Windows Multi-Head Self-Attention—(S)W-MSA的概念，
  • 比如在下图的4倍下采样和8倍下采样中，将特征图划分成了多个不相交的区域（Window），并且Multi-Head Self-Attention只在每个窗口（Window）内进行。相对于Vision Transformer中直接对整个（Global）特征图进行Multi-Head Self-Attention，这样做的目的是能够减少计算量的，尤其是在浅层特征图很大的时候。
  • 这样做虽然减少了计算量但也会隔绝不同窗口之间的信息传递，所以在论文中作者又提出了 Shifted Windows Multi-Head Self-Attention(SW-MSA)的概念，通过此方法能够让信息在相邻的窗口中进行传递，后面会讲解。
    

3. 模型详解

接下来，我们看下原论文中给出的关于Swin Transformer（Swin-T）网络的架构图。通过图 (a) 可以看出整个框架的基本流程如下：

• 首先，将图片输入到 Patch Partition 模块中进行分块，即每4x4相邻的像素为一个Patch，然后在channel方向展平（flatten）。
  • 假设输入的是RGB三通道图片，那么每个patch就有4x4=16个像素，然后每个像素有R、G、B三个值所以展平后是16x3=48，所以通过Patch Partition后图像shape由 [H, W, 3]变成了 [H/4, W/4, 48]。
  • 然后在通过 Linear Embeding 层对每个像素的channel数据做线性变换，由48变成 C，即图像shape再由 [H/4, W/4, 48] 变成了 [H/4, W/4, C]。
  • 其实在源码中Patch Partition和Linear Embeding就是直接通过一个卷积层实现的，和之前 Vision Transformer中讲的 Embedding层结构一模一样。
• 然后，就是通过四个Stage构建不同大小的特征图，
  • 除了Stage1中先通过一个Linear Embeding层外，
  • 剩下三个stage都是先通过一个Patch Merging层进行下采样（后面会细讲）。
  • 之后，都是重复堆叠Swin Transformer Block注意这里的Block其实有两种结构，如图(b)中所示，这两种结构的不同之处仅在于一个使用了W-MSA结构，一个使用了SW-MSA结构。而且这两个结构是成对使用的，先使用一个W-MSA结构再使用一个SW-MSA结构。所以你会发现堆叠Swin Transformer Block的次数都是偶数（因为成对使用）。
• 最后，对于分类网络，后面还会接上一个Layer Norm层、全局池化层以及全连接层得到最终输出。图中没有画，但源码中是这样做的。

接下来，在分别对Patch Merging、W-MSA、SW-MSA以及使用到的相对位置偏执（relative position bias）进行详解。关于Swin Transformer Block中的MLP结构和Vision Transformer中的结构是一样的，可以参考：Vision Transformer。

3.1 Patch Merging

前面有说，在所有的Stage中首先要通过一个 Patch Merging 层进行下采样（Stage1除外）。如下图所示，假设输入Patch Merging的是一个4x4大小的单通道特征图（feature map），

• 首先，Patch Merging会将每个2x2的相邻像素划分为一个patch，
• 然后将每个patch中相同位置（同一颜色）像素给拼在一起就得到了4个feature map。
• 接着将这四个feature map在深度方向进行concat拼接，然后在通过一个LayerNorm层。
• 最后通过一个全连接层在feature map的深度方向做线性变化，将feature map的深度由C变成C/2。

通过这个简单的例子可以看出，通过Patch Merging层后，feature map的高和宽会减半，深度会翻倍。

3.2 Windows Multi-head Self-Attention（W-MSA）

引入Windows Multi-head Self-Attention（W-MSA）模块是为了减少计算量。如下图所示，

• 左侧使用的是普通的Multi-head Self-Attention（MSA）模块，对于feature map中的每个像素（或称作token，patch）在Self-Attention计算过程中需要和所有的像素去计算。
• 但在图右侧，在使用Windows Multi-head Self-Attention（W-MSA）模块时，首先将feature map按照MxM（例子中的M=2）大小划分成一个个Windows，然后单独对每个Windows内部进行Self-Attention。

两者的计算量具体差多少呢？原论文中有给出下面两个公式，这里忽略了Softmax的计算复杂度。

• h代表feature map的高度
• w代表feature map的宽度
• C代表feature map的深度
• M代表每个窗口（Windows）的大小

这个公式是的计算原论文中并没有细讲，这里解释一下。Self-Attention的理论需要先知道，可以参考：Transformer。

3.2.1 MSA（Multi-head Self-Attention） 计算量

对于feature map中的每个像素（或称作token，patch），都要通过 $W_q,W_k,W_v$ 生成对应的$query(q)$，$key(k)$ 以及 $value(v)$。这里假设 $q,k,v$ 的向量长度与feature map的深度 $C$ 保持一致。那么对应所有像素生成Q的过程如下式：

根据矩阵运算的计算量公式可以得到生成Q的计算量为 $hw\times C\times C$，生成 K 和 V 同理都是$hw{C}^2$ ，那么总共是 $3hw{C}^2$ 。接下来 Q 和 $K^T$ 相乘，对应计算量为 $(hw{)}^{2}C$：

接下来忽略除以 $\sqrt{d}$，以及softmax的计算量，假设得到 ${\Lambda }^{hw\times hw}$ ，最后还要乘以 $V$，对应的计算量为 $(hw{)}^2$：

那么对应单头的Self-Attention模块，总共需要 $3hw{C}^2+(hw{)}^{2}C+(hw{)}^{2}C=3hw{C}^2+2(hw{)}^{2}C$。而在实际使用过程中，使用的是多头的Multi-head Self-Attention模块，在之前的文章中有进行过实验对比，多头注意力模块相比单头注意力模块的计算量仅多了最后一个融合矩阵 $W_O$ 的计算量$hw{C}^2$。

因此，上述计算量总共加起来是： $4hw{C}^2+2(hw{)}^{2}C$。

3.2.2 W-MSA（Windows Multi-head Self-Attention）计算量

对于W-MSA模块首先要将feature map划分到一个个窗口（Windows）中，假设每个窗口的宽高都是M，那么总共会得到 $\frac{h}{M}\times \frac{w}{M}$个窗口，然后对每个窗口内使用多头注意力模块。刚刚计算高为 h，宽为 w，深度为 C 的feature map的计算量为：$4hw{C}^2+2(hw{)}^{2}C$，这里每个窗口的高为 M 宽为 M，带入公式得：

又因为有 $\frac{h}{M}\times \frac{w}{M}$ 个窗口，则：

故使用W-MSA模块的计算量为： $4hw{C}^2+2M^{2}hwC$。

假设feature map的h、w都为112，M=7，C=128，采用W-MSA模块相比MSA模块能够节省约40124743680 FLOPs：

3.3 Shifted Windows Multi-Head Self-Attention（ SW-MSA）

• 由于采用W-MSA模块，只会在每个窗口内进行自注意力计算，所以窗口与窗口之间是无法进行信息传递的。

为了解决这个问题，作者引入了Shifted Windows Multi-Head Self-Attention（SW-MSA）模块，即进行偏移的 W-MSA。如下图所示，

• 左侧使用的是刚刚讲的W-MSA（假设是第L层），根据之前介绍的W-MSA和SW-MSA是成对使用的，那么第L+1层使用的就是SW-MSA（右侧图）。
• 根据左右两幅图对比能够发现：窗口（Windows）发生了偏移（可以理解成窗口从左上角分别向右侧和下方各偏移了 $\lfloor \frac{M}{2}\rfloor$ 个像素）。
  • 偏移后的窗口（右侧图）中，比如对于第一行第2列的2x4的窗口，它能够使第L层的第一排的两个窗口信息进行交流。再比如，第二行第二列的4x4的窗口，他能够使第L层的四个窗口信息进行交流，其他的同理。那么这就解决了不同窗口之间无法进行信息交流的问题。

上图中，我们可以发现通过将窗口进行偏移后，由原来的4个窗口变成9个窗口了。后面又要对每个窗口内部进行MSA，这样做感觉又变麻烦了。

• 为了解决这个麻烦，作者又提出而了 Efficient batch computation for shifted configuration，一种更加高效的计算方法。下面是原论文给的示意图：

上图有些难以理解，参考：太阳花重新画的。下图左侧是刚刚通过偏移窗口后得到的新窗口，右侧是为了方便大家理解，对每个窗口加上了一个标识。

• 0对应的窗口标记为区域A，3和6对应的窗口标记为区域B，1和2对应的窗口标记为区域C。
  
• 先将区域A和C移到最下方：
  
• 接着，再将区域A和B移至最右侧：
  

移动完后，如上图所示：

• 4是一个单独的窗口；
• 将5和3合并成一个窗口；
• 7和1合并成一个窗口；
• 8、6、2和0合并成一个窗口。

这样又和原来一样是4个4x4的窗口了，所以能够保证计算量是一样的。这里肯定有人会有疑问，把不同的区域合并在一起（比如5和3）进行MSA，这信息不就乱窜了吗？

• 确实会出现问题，为了防止这个问题，在实际计算中使用的是masked MSA即带蒙板mask的MSA，这样就能够通过设置蒙板来隔绝不同区域的信息了。关于mask如何使用，参考下面这幅图，下图是以上面的区域 5 和3 为例。
  

对于该窗口内的每一个像素（或称token，patch）在进行MSA计算时，都要先生成对应的query(q)，key(k)，value(v)。假设对于上图的像素0而言，得到 $q^0$ 后要与每一个像素的 k 进行匹配（match），假设 ${\alpha }_{0,0}$ 代表 $q^0$ 与像素 0 对应的 $k^0$ 进行匹配的结果，那么同理可以得到 ${\alpha }_{0,0}$ 至 ${\alpha }_{0,15}$。按照普通的MSA计算，接下来就是 SoftMax 操作了。

• 但对于这里的masked MSA，像素0是属于区域5的，我们只想让它和区域5内的像素进行匹配。
• 那么我们可以将像素0与区域3中的所有像素匹配结果都减去100（例如 ${\alpha }_{0,2},{\alpha }_{0,3},{\alpha }_{0,6},{\alpha }_{0,7}$等等）。
• 由于 $\alpha$ 的值都很小，一般都是零点几的数，将其中一些数减去100后在通过SoftMax得到对应的权重都等于0了。
• 所以对于像素 0 而言实际上还是只和区域5内的像素进行了MSA。那么对于其他像素也是同理。

具体代码是怎么实现的，可以参考：其代码解析。

• 注意，在计算完后还要把数据给挪回到原来的位置上（例如上述的A，B，C区域）。

3.4 Relative Position Bias

关于相对位置偏执，论文里也没有细讲，只说了参考的哪些论文，然后说使用了相对位置偏执后给够带来明显的提升。

• 根据原论文中的表4可以看出，在Imagenet数据集上如果不使用任何位置偏执，top-1为80.1，但使用了相对位置偏执（rel. pos.）后top-1为83.3，提升还是很明显的。
  

1）这个相对位置偏执是加在哪的呢？

• 根据论文中提供的公式可知是在Q和K进行匹配并除以 d \\sqrt d

  Swin-Transformer网络结构详解

  文章目录

  • 0 前言
  • 1 网络整体框架
  • 2 Patch Merging详解
  • 3 W-MSA详解
  • • MSA模块计算量
    • W-MSA模块计算量
  • 4 SW-MSA详解
  • 5 Relative Position Bias详解
  • 6 模型详细配置参数

  ---

  0 前言

  Swin Transformer是2021年微软研究院发表在ICCV上的一篇文章，并且已经获得ICCV 2021 best paper的荣誉称号。Swin Transformer网络是Transformer模型在视觉领域的又一次碰撞。该论文一经发表就已在多项视觉任务中霸榜。该论文是在2021年3月发表的，现在是2021年11月了，根据官方提供的信息可以看到，现在还在COCO数据集的目标检测以及实例分割任务中是第一名（见下图State of the Art表示第一）。

  

  论文名称：Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
  原论文地址： https://arxiv.org/abs/2103.14030
  官方开源代码地址：https://github.com/microsoft/Swin-Transformer

  Pytorch实现代码： pytorch_classification/swin_transformer
  Tensorflow2实现代码：tensorflow_classification/swin_transformer

  不想看文章的可以看下我在bilibili上讲的视频： https://www.bilibili.com/video/BV1pL4y1v7jC

  ---

  1 网络整体框架

  在正文开始之前，先来简单对比下Swin Transformer和之前的Vision Transformer（如果不了解Vision Transformer的建议先去看下我之前的文章）。下图是Swin Transformer文章中给出的图1，左边是本文要讲的Swin Transformer，右边边是之前讲的Vision Transformer。通过对比至少可以看出两点不同：

  • Swin Transformer使用了类似卷积神经网络中的层次化构建方法（Hierarchical feature maps），比如特征图尺寸中有对图像下采样4倍的，8倍的以及16倍的，这样的backbone有助于在此基础上构建目标检测，实例分割等任务。而在之前的Vision Transformer中是一开始就直接下采样16倍，后面的特征图也是维持这个下采样率不变。
  • 在Swin Transformer中使用了Windows Multi-Head Self-Attention(W-MSA)的概念，比如在下图的4倍下采样和8倍下采样中，将特征图划分成了多个不相交的区域（Window），并且Multi-Head Self-Attention只在每个窗口（Window）内进行。相对于Vision Transformer中直接对整个（Global）特征图进行Multi-Head Self-Attention，这样做的目的是能够减少计算量的，尤其是在浅层特征图很大的时候。这样做虽然减少了计算量但也会隔绝不同窗口之间的信息传递，所以在论文中作者又提出了 Shifted Windows Multi-Head Self-Attention(SW-MSA)的概念，通过此方法能够让信息在相邻的窗口中进行传递，后面会细讲。

  

  接下来，简单看下原论文中给出的关于Swin Transformer（Swin-T）网络的架构图。通过图(a)可以看出整个框架的基本流程如下：

  

  • 首先将图片输入到Patch Partition模块中进行分块，即每4x4相邻的像素为一个Patch，然后在channel方向展平（flatten）。假设输入的是RGB三通道图片，那么每个patch就有4x4=16个像素，然后每个像素有R、G、B三个值所以展平后是16x3=48，所以通过Patch Partition后图像shape由 [H, W, 3]变成了 [H/4, W/4, 48]。然后在通过Linear Embeding层对每个像素的channel数据做线性变换，由48变成C，即图像shape再由 [H/4, W/4, 48]变成了 [H/4, W/4, C]。其实在源码中Patch Partition和Linear Embeding就是直接通过一个卷积层实现的，和之前Vision Transformer中讲的 Embedding层结构一模一样。

  • 然后就是通过四个Stage构建不同大小的特征图，除了Stage1中先通过一个Linear Embeding层外，剩下三个stage都是先通过一个Patch Merging层进行下采样（后面会细讲）。然后都是重复堆叠Swin Transformer Block注意这里的Block其实有两种结构，如图(b)中所示，这两种结构的不同之处仅在于一个使用了W-MSA结构，一个使用了SW-MSA结构。而且这两个结构是成对使用的，先使用一个W-MSA结构再使用一个SW-MSA结构。所以你会发现堆叠Swin Transformer Block的次数都是偶数（因为成对使用）。

  • 最后对于分类网络，后面还会接上一个Layer Norm层、全局池化层以及全连接层得到最终输出。图中没有画，但源码中是这样做的。

  接下来，在分别对Patch Merging、W-MSA、SW-MSA以及使用到的相对位置偏执（relative position bias）进行详解。关于Swin Transformer Block中的MLP结构和Vision Transformer中的结构是一样的，所以这里也不在赘述，参考。

  ---

  2 Patch Merging详解

  前面有说，在每个Stage中首先要通过一个Patch Merging层进行下采样（Stage1除外）。如下图所示，假设输入Patch Merging的是一个4x4大小的单通道特征图（feature map），Patch Merging会将每个2x2的相邻像素划分为一个patch，然后将每个patch中相同位置（同一颜色）像素给拼在一起就得到了4个feature map。接着将这四个feature map在深度方向进行concat拼接，然后在通过一个LayerNorm层。最后通过一个全连接层在feature map的深度方向做线性变化，将feature map的深度由C变成C/2。通过这个简单的例子可以看出，通过Patch Merging层后，feature map的高和宽会减半，深度会翻倍。

  

  ---

  3 W-MSA详解

  引入Windows Multi-head Self-Attention（W-MSA）模块是为了减少计算量。如下图所示，左侧使用的是普通的Multi-head Self-Attention（MSA）模块，对于feature map中的每个像素（或称作token，patch）在Self-Attention计算过程中需要和所有的像素去计算。但在图右侧，在使用Windows Multi-head Self-Attention（W-MSA）模块时，首先将feature map按照MxM（例子中的M=2）大小划分成一个个Windows，然后单独对每个Windows内部进行Self-Attention。

  
  两者的计算量具体差多少呢？原论文中有给出下面两个公式，这里忽略了Softmax的计算复杂度。：
  Ω ( M S A ) = 4 h w C 2 + 2 ( h w ) 2 C                 ( 1 ) Ω ( W − M S A ) = 4 h w C 2 + 2 M 2 h w C       ( 2 ) \\Omega (MSA)=4hwC^2 + 2(hw)^2C \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ (1) \\\\ \\Omega (W-MSA)=4hwC^2 + 2M^2hwC\\ \\ \\ \\ \\ (2) Ω(MSA)=4hwC2+2(hw)2C               (1)Ω(W−MSA)=4hwC2+2M2hwC     (2)

  • h代表feature map的高度
  • w代表feature map的宽度
  • C代表feature map的深度
  • M代表每个窗口（Windows）的大小

  这个公式是咋来的，原论文中并没有细讲，这里简单说下。首先回忆下单头Self-Attention的公式，如果对Self-Attention不了解的，请看下我之前写的文章。
  A t t e n t i o n ( Q , K , V ) = S o f t M a x ( Q K T d ) V Attention(Q, K, V)=\\rm SoftMax(\\fracQK^T\\sqrt d)V Attention(Q,K,V)=SoftMax(d ​QKT​)V

  MSA模块计算量

  对于feature map中的每个像素（或称作token，patch），都要通过 W q , W k , W v W_q, W_k, W_v Wq​,Wk​,Wv​生成对应的query(q)，key(k)以及value(v)。这里假设q, k, v的向量长度与feature map的深度C保持一致。那么对应所有像素生成Q的过程如下式：
  A h w × C ⋅ W q C × C = Q h w × C A^hw \\times C \\cdot W^C \\times C_q=Q^hw \\times C Ahw×C⋅WqC×C​=Qhw×C

  • A h w × C A^hw \\times C Ahw×C为将所有像素（token）拼接在一起得到的矩阵（一共有hw个像素，每个像素的深度为C）
  • W q C × C W^C \\times C_q WqC×C​为生成query的变换矩阵
  • Q h w × C Q^hw \\times C Qhw×C为所有像素通过 W q C × C W^C \\times C_q WqC×C​得到的query拼接后的矩阵

  根据矩阵运算的计算量公式可以得到生成Q的计算量为 h w × C × C hw \\times C \\times C hw×C×C，生成K和V同理都是 h w C 2 hwC^2 hwC2，那么总共是 3 h w C 2 3hwC^2 3hwC2。接下来 Q Q Q和 K T K^T KT相乘，对应计算量为 ( h w ) 2 C (hw)^2C (hw)2C：
  Q h w × C ⋅ K T ( C × h w ) = X h w × h w Q^hw \\times C \\cdot K^T(C \\times hw)= X^hw \\times hw Qhw×C⋅KT(C×hw)=Xhw×hw
  接下来忽略除以 d \\sqrt d d ​以及softmax的计算量，假设得到 Λ h w × h w \\Lambda ^hw \\times hw Λhw×hw，最后还要乘以V，对应的计算量为 ( h w ) 2 C (hw)^2C (hw)2C:
  Λ h w × h w ⋅ V h w × C = B h w × C \\Lambda ^hw \\times hw \\cdot V^hw \\times C=B^hw \\times C Λhw×hw⋅Vhw×C=Bhw×C
  那么对应单头的Self-Attention模块，总共需要 3 h w C 2 + ( h w ) 2 C + ( h w ) 2 C = 3 h w C 2 + 2 ( h w ) 2 C 3hwC^2 + (hw)^2C + (hw)^2C=3hwC^2 + 2(hw)^2C 3hwC2+(hw)2C+(hw)2C=3hwC2+2(hw)2C。而在实际使用过程中，使用的是多头的Multi-head Self-Attention模块，在之前的文章中有进行过实验对比，多头注意力模块相比单头注意力模块的计算量仅多了最后一个融合矩阵 W O W_O WO​的计算量 h w C 2 hwC^2 hwC2。
  B h w × C ⋅ W O C × C = O h w × C B^hw \\times C \\cdot W_O^C \\times C = O^hw \\times C Bhw×C⋅WOC×C​=Ohw×C
  所以总共加起来是： 4 h w C 2 + 2 ( h w ) 2 C 4hwC^2 + 2(hw)^2C 改进YOLOv5系列：增加Swin-Transformer小目标检测头

  计算机视觉方向简介 | 视觉惯性里程计(VIO)

  pytorch 笔记： Swin-Transformer 代码

  Swin Transformer对CNN的降维打击

  Swin Transformer对CNN的降维打击

  Swin Transformer对CNN的降维打击