Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer

论文链接:Swin Transformer

论文代码:https://github.com/microsoft/Swin-Transformer

目录

1、摘要和背景介绍

2、整体框架

2.1、基于移位窗口的自注意力

2.2、移动窗口划分模式

2.3、两个连续的Swin transformer块

2.4、Swin的几种变体

4、实验部分

5、结论


1、摘要和背景介绍

        Transformer在NLP领域取得不错表现,如何更好地处理图像成为行业所面临的问题。图像中的像素与文本中的单词相比分辨率较高,这就给从语言到视觉的转换带来了挑战。

        为了解决这些差异,本文提出了一种分层Transformer,它的表示是用shifted windows(移位窗口方案)计算的。移位窗口方案通过将 self-attention 计算限制在不重叠的局部窗口上,同时还允许跨窗口连接,从而带来更高的效率。

        这种分层架构在各种尺度上使用较为灵活,并且对于图像大小具有线性计算复杂性。

        在许多视觉任务,例如语义分割,需要在像素级别进行密集预测, Transformer 在高分辨率图像上的处理较为困难,因为自注意力的计算复杂度与图像大小成二次方。

        ViT中是将图像分割成多个块,将每个图像块作为一个token,并没有考虑图像的三个通道,而是从全局特征的角度考虑整幅图像,如果此时遇到高分辨率的图像,采用ViT处理会产生极大的计算复杂度,此时的计算复杂度与图像大小成二次方。

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第1张图片

可以参考我之前的一篇博客:Transformer用于大规模图像识别

        本文提出了一种通用的骨干网络,称为 Swin Transformer,它构建分层特征图,并且对图像大小具有线性计算复杂度。

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第2张图片

        上图 (a) 表示Swin Transformer 通过合并更深层的图像块(以灰色显示)来构建分层特征图,并且由于仅在每个局部窗口内计算自注意力,因此对输入图像大小具有线性计算复杂度 红色的)。  图(b) 是ViT生成单个低分辨率的特征图,并且由于全局自注意力的计算,输入图像大小具有二次方计算复杂度。

2、整体框架

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第3张图片

         上图展示了 Swin Transformer 的整体架构 (Swin-T)。 首先通过像 ViT 的图像分割模块将输入的 RGB 图像分割成不重叠的patch。 每个patch都被看作一个“token”,其特征设置为原始 RGB图像像素值的串联。 本文采用 4 × 4 的patch大小,因此每个patch的特征维度为 4 × 4 × 3 = 48(这里的3表示图像的通道数)。然后采用投影将特征映射到任意维度( 记为 C)。

这部分就对应上图中的patch partition和linear Embedding,可以参考ViT中的实现方法来理解,只是这里考虑了通道数。

        在"stage1"中,有两个Swin Transformer块,结构如图(b)所示,有两种改进的自注意力transformer:W-MSA和SW-MSA,分别表示常规的MSA和移位窗口的MSA。Transformer 块保持token的数量(H/4 ×W/4)。

        为了产生分层表示,随着网络变得更深,文中通过patch merging减少token的数量。

        第一个patch merging层连接每组 2×2 相邻patch的特征,并在 4C 维拼接特征上应用线性层。 这将token的数量减少了 2 × 2 = 4 的倍数(2× 分辨率(4×4)的下采样),并且输出维度设置为 2C。 之后应用 Swin Transformer 块进行特征变换,分辨率保持在 H/8 ×W/8 。第一个块合并(patch merging)和特征转换被称为“stage 2”

什么是patch merging?

Patch merge层通过拼接2×2个相邻Patch的特征来减少token的数量,并在拼接的4C维特征上应用线性层使隐藏层的维数翻倍。这个操作类似于CNN的空洞卷积。

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第4张图片

参考博客: 图解Swin Transformer - 搜索结果 - 知乎

        重复stage 2两次,分别为“Stage 3”和“Stage 4”,输出分辨率分别为 H/16 × W/16 和 H/32 × W/32 。这些阶段共同产生分层表示。类似于VGG和ResNet网络相同的特征图分辨率。

        对于上图(b)的两种Swin Transformer块,通过将 Transformer 块中的标准多头自注意力 (MSA) 模块替换为基于移位窗口(shifted windows)的模块构建的,其他层保持不变。一个 Swin Transformer 模块由一个基于移动窗口的 MSA 模块组成,然后是一个 2 层 MLP,其间具有 GELU 非线性。 在每个 MSA 模块和每个 MLP 之前应用一个 LayerNorm (LN) 层,在每个模块之后应用一个残差连接。

对于残差连接的理解可以参考之前的博客:https://blog.csdn.net/weixin_43687366/article/details/123316333?spm=1001.2014.3001.5501

最开始的Transformer中的MSA:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第5张图片

2.1、基于移位窗口的自注意力

        ViT从全局窗口角度计算自注意力,计算量比较大,本文提出在局部窗口内计算自注意力,窗口以不重叠的方式均匀地划分。假设每个窗口包含 M × M 个patch,全局 MSA 模块和基于 h × w 个patch图像的窗口的计算复杂度为 :

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第6张图片

         前者的计算复杂度为 hw 的二次方,后者在 M 固定时是线性的(默认设置为 7)。 全局自注意力相比基于窗口的自注意力而言计算复杂度高且可扩展性低。

这里的h和w分别表示原始图像的高和宽,C表示通道数,M表示每个图像块的大小,对应图像块的h和w均为M;

复杂度计算过程:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第7张图片

 首先,to_q_k_v()函数,对于一个hw*c的特征矩阵,对于Q而言是用一个h*w*c的向量乘以一个c*c的系数矩阵,得到了h*w*c,所以每一个计算的复杂度是h*w*c^2,因为总计有三次操作(Q、K、V),因此是3倍的hwc^2。

然后,自注意力的计算是h*w*c乘以k的转置(c*h*w),所以得到h*w*h*w=(hw)^2,这个计算复杂度就是(h*w)^2*c。

接下来,自注意力矩阵和value的乘积的计算复杂度还是(h*w)^2*c,所以现在就是2*(h*w)^2*c;

还有最后一步,映射层是h*w*c乘以c*c变成了h*w*c,这里的计算复杂度是h*w*c^2;

最后将上述的计算复杂度加在一起就是:4hwC^2+2(hw)^2*C,即公式(1)所示。

公式(1)中的计算复杂度是二次方级别的,计算较为费时,本文提出了一种线性的方式,提升计算速度,如公式(2)所示:

基于窗口的自注意力计算复杂度:

这里还是在每个窗口中计算多头注意力,直接套用公式(1),只不过高度和宽度发生变化,现在的窗口变成了M*M,而不是h和w,将M的值代入公式(1)中就得到4*M^2*C^2+2*M^4*C,这表示在一个窗口里计算多头注意力所需的计算复杂度。

现在一共有h/M*w/M个窗口,将这些窗口乘以每个窗口所需要的计算复杂度就得到公式(2)。看着差别不大,可以尝试代入数字,h*w=50*50=2500,而M^2=49,这里相差了2500/49=50倍,甚至是上百倍。所以这种窗口的自注意力方式可以很好地解决内存和计算量的问题。但是这种方式缺少窗口和窗口之间的通信,达不到全局建模的目的,于是本文又提出了一种移动窗口的方式,可以让窗口和窗口之间进行通信,并且具有上下文的信息。

相对位置偏置:

 主要是在原始计算Attention的公式中的Q,K时加入了相对位置编码,后续实现也证明了相对位置编码的加入提升了模型的性能,如表4所示:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第8张图片

2.2、移动窗口划分模式

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第9张图片

         Layer 1表示常规的窗口划分方式,类似于ViT中的划分,特征图有4个窗口(每个窗口又是4*4个小块),但是做完移动窗口(即layer 1+1)之后,得到9个窗口,窗口的数量增加了,而且每个窗口的里的元素大小不一,中间的窗口是4*4,然后还有2*2和4*2类型的窗口。

        layer1+1的这种形式很难计算自注意力,又或者将上下左右的块都通过pad填充成中间的patch大小,表示成3*3的patch,但是这种方式的计算复杂度比较高。

        因此,为了保证移位之后的窗口数量还是4个,并且每个patch中的大小还保持一致(4*4=16),于是文中就提出了一种掩码的方式,如下图4所示:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第10张图片

         上图 window partition 部分表示通过普通的窗口移动方式得到的9个窗口,并不是直接对这9个窗口计算自注意力,而是通过一次循环移位(cyclic shift)操作

即原始的A和C直接移位到最下面一行,而原始的B直接移位到最右边,得到新的移位窗口,这个新的移位窗口的数量就为4,每个patch中包括4*4=16个小块,这样的话窗口的数量固定了,计算复杂度也就固定了。

        这种移位窗口的方式,使得相邻的组(patch)进行信息交互,这个有点类似于ShuffleNet,但这个是在空间维度上的shuffle,而ShuffleNet是在通道维度上的操作。

        但是新的问题出现了,新的移位窗口中,左上角的区域可以很方便的计算自注意力,但是其他块怎么计算呢?因为它们当中的元素都是从其他地方移过来的,所以不用计算自注意力,因为彼此元素块在图像中并没有太大的联系。因此本文采用一种掩码的方式计算自注意力,即图4中的masked MSA。

        做完掩码操作之后,又将移位的窗口还原回去,即还原成window partition的形式。主要是为了保持图像的相对位置,整体图像的语义信息也是不变的。

掩码操作的具体思路?

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第11张图片

对移位后的窗口,划分为4个窗口即上图的window0、window1、window2和window3, 对于Window2中的区域3和区域6的计算过程如下图所示:对于36和63不需要计算注意力的区域通过设置一个掩码,与上述注意力矩阵进行相加,而36和63中的值比较小,经过softmax之后变为0。

其他区域计算同样按照下面的方式:

作者通过这种循环位移的方式和巧妙的掩码模板,从而实现了一次前向过程就能把所需的自注意力值都计算出来,而且只需要计算4个窗口,也就是说窗口的数量没有增加,计算复杂度也没有增加,就完成了这个任务。

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第12张图片

2.3、两个连续的Swin transformer块

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第13张图片

 这里的连续Swin Transformer块是先进行window,然后进行shifted window,可以起到窗口和窗口之间的通信,每次都是先经历一次基于窗口的多头自注意力,然后再经历一次基于移动窗口的多头自注意力。具体过程如上述所示:

每次输入进来先经过一次LayerNorm,然后就是窗口的多头自注意力,然后再通过一个LayerNorm和MLP,第一个Block完成;

然后经过一个shifted window,也就是基于移动窗口的多头自注意力,然后通过LayerNorm和MLP得到输出。

这两个块加起来才是Swin Transformer的基本的计算单元,这也就是为什么图1中的每个stage都是偶数,因为始终都需要两层block连在一起作为一个基本单元,所以一定是2的倍数。

2.4、Swin的几种变体

        本文的基础模型是 Swin-B,具有类似于 ViT-B/DeiT-B 的模型大小和计算复杂度。 还提出了 Swin-T、Swin-S 和 Swin-L,它们分别是模型大小和计算复杂度的 0.25 倍、0.5 倍和 2 倍的版本。 其中Swin-T 和 Swin-S 的复杂度分别类似于 ResNet-50 (DeiT-S) 和 ResNet-101。 默认情况下,窗口大小设置为 M = 7。 对于所有实验,每个头部的查询维度为 d = 32,每个 MLP 的扩展层为 α = 4。 这些模型变体的架构超参数是:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第14张图片

 可以看到这些变体之间主要不同在于:

1、一是向量维度的大小为C;

2、每个stage中所包含的transformer block的块数不一样;

比较像ResNet的设计结构:ResNet也是分了4个stage,每个stage中也有不同数量的残差块;

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第15张图片

Swin Transformer:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第16张图片

4、实验部分

        对 ImageNet-1K 图像分类 、COCO 目标检测 和 ADE20K 语义分割进行了实验。 并将提出的 Swin Transformer 架构与之前在三个任务上的最新技术进行比较。 然后做了一些消融实验。

图像分类:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第17张图片

 目标检测:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第18张图片

 语义分割:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第19张图片

 消融实验:

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第20张图片

 不同的自注意力方法的计算速度对比

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第21张图片

Swin-Transformer:基于移位窗口(Shifted Windows)的分层视觉Transformer_第22张图片

 5、结论

        本文介绍了 Swin Transformer,这是一种新的视觉 Transformer,可产生分层特征表示和相对于输入图像大小具有线性计算复杂度。Swin Transformer 在 COCO 目标检测和 ADE20K 语义分割方面实现了最先进的性能,超越了以前的最佳方法。 

        作为 Swin Transformer 的一个关键元素,基于移位窗口的自注意力在视觉问题上被证明是有效且高效的。

参考视频:Swin Transformer论文精度

你可能感兴趣的:(论文翻译,transformer,深度学习,人工智能,SwinTransformer)