[深度学习论文笔记]Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

论文：
Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation
Swin-Unet：用于医学图像分割的类Unet纯transformer
Published: 2021 May
论文：https://arxiv.org/abs/2105.05537
代码：https://github.com/HuCaoFighting/Swin-Unet

问题动机：
得益于深度学习的发展，计算机视觉技术在医学图像分析中得到了广泛的应用。图像分割是医学图像分析的重要组成部分。特别是准确、鲁棒的医学图像分割可以在计算机辅助诊断和图像引导的临床手术中发挥基石作用。
现有的医学图像分割方法主要依赖u型结构的全卷积神经网(FCNN)。典型的u型网络，U-Net，由一个具有跳跃连接的对称编码器-解码器组成。在编码器中，采用一系列卷积层和连续降采样层来提取接收感受野的深度特征。然后，解码器将提取的深度特征向上采样到输入分辨率进行像素级语义预测，并通过跳跃连接融合来自编码器的不同尺度的高分辨率特征，以减轻降采样导致的空间信息丢失。由于这种优雅的结构设计，U-Net在各种医学成像应用中取得了巨大的成功。遵循这一技术路线，许多算法如3D U-Ne、Res-UNet、U-Net++和UNet3+被开发出来用于各种医学成像模式的图像和体积分割。这些基于fcnn的方法在心脏分割、器官分割和病变分割方面的优异性能证明了CNN具有很强的学习判别特征的能力。

思路来源：
目前，虽然基于CNN的方法在医学图像分割领域取得了优异的性能，但仍不能完全满足医学应用对分割精度的严格要求。在医学图像分析中，图像分割仍然是一个具有挑战性的课题。由于卷积运算固有的局限性，基于CNN的方法很难学习显式的全局和长期语义信息交互。一些研究试图通过使用深度卷积层、自我注意机制和图像金字塔来解决这个问题。然而，这些方法在建模长期依赖时仍然有局限性。然而受Transformer在自然语言处理(NLP)领域的巨大成功的启发，研究人员试图将Transformer引入视觉领域。在[An image is worth 16x16 words: Transformers for image recognition at scale]中，视觉transformer(vision transformer, ViT)被提出来执行图像识别任务。以具有位置嵌入的二维图像块为输入，在大数据集上进行预处理，与基于CNN的方法取得了相当的性能。此外，[Training data-efficient image transformers & distillation through attention]中还提出了数据高效的图像变换(DeiT)，这表明transformer可以在中等规模的数据集上进行训练，并将其与蒸馏方法相结合，可以得到更鲁棒的transformer。在[Swin Transformer: Hierarchical Vision Transformer using Shifted Windows]中，作者开发了一个分层的Swin transformer。作者以Swin Transformer为视觉中枢，在图像分类、目标检测和语义分割方面取得了最先进的性能。ViT、DeiT和Swin transformer在图像识别领域的成功证明了transformer在视觉领域的应用潜力。
在Swin Transformer成功的激励下，作者提出Swin- unet来利用Transformer在2D医学图像分割中的威力。swin-unet是第一个纯粹的基于transformer的u型架构，它由编码器、瓶颈、解码器和跳跃连接组成。编码器、瓶颈和解码器都是基于Swin-transformer模块构建的。将输入的医学图像分割成不重叠的图像补丁。每个补丁都被视为一个token，并被输入到基于transformer的编码器中，以学习深度特征表示。提取的上下文特征由带补丁扩展层的解码器上采样，通过跳跃连接与编码器的多尺度特征融合，恢复特征图的空间分辨率，进一步进行分割预测。在多器官和心脏数据集上的大量实验表明，该方法具有良好的分割精度和鲁棒泛化能力。
具体来说，作者的贡献可以总结为:
(1)基于Swin Transformer块，构建了一个具有跳跃连接的对称编码器-解码器体系结构。在编码器中实现了从局部到全局的自注意;在解码器中，将全局特征上采样到输入分辨率，进行相应的像素级分割预测。
(2)开发了patch扩展层，无需卷积或插值操作即可实现上采样和特征维数的增加。
(3)实验发现跳跃连接对transformer也是有效的，因此最终构建了一个纯基于transformer的u型编解码结构，具有跳跃连接，命名为swin-unet。

主要结构:
所提出的swin-unet的整体架构如图1所示。swin - unet由编码器、瓶颈、解码器和跳过连接组成。swin - unet的基本单元是Swin transformer块。对于编码器，为了将输入转换为序列嵌入，将医学图像分割成大小为4 × 4的非重叠patch。通过这种划分方法，每个patch的特征维数变为4 × 4 × 3 = 48。此外，将投影的特征维度应用线性嵌入层(以C表示)，转换后的补丁标记经过几个Swin Transformer块和补丁合并层生成分层特征表示。其中，补丁合并层负责下采样和增加维度，Swin Transformer块负责特征表示学习，受U-Net的启发，设计了一种基于对称transformer的解码器。该解码器由Swin transformer块和补丁扩展层组成。提取的上下文特征通过跳跃连接与编码器的多尺度特征融合，以弥补降采样造成的空间信息丢失。与补丁合并层不同，补丁扩展层被专门设计用于执行上采样。补丁扩展层通过2×上采样分辨率将相邻维度的特征地图重塑为一个大的特征地图。最后，利用最后一个补丁扩展层进行4×上采样，将特征映射的分辨率恢复到输入分辨率(W ×H)，然后在这些上采样特征上应用线性投影层输出像素级分割预测。

图一：swin-unet的体系结构由编码器、瓶颈、解码器和跳过连接组成。编码器、瓶颈和解码器都是基于swin transformer块构造的

结构详解：
Swin Transformer块：

与传统的多头自注意(MSA)模块不同，swin transformer块是基于平移窗口构造的。在图2中，给出了两个连续的swin transformer块。每个swin transformer块由LayerNorm (LN)层、多头自注意模块、剩余连接和具有GELU非线性的2层MLP组成。在两个连续的transformer模块中分别采用了基于窗口的多头自注意(W-MSA)模块和位移的基于窗口的多头自注意(SW-MSA)模块。基于这种窗口划分机制，连续swin transformer块可表示为:

自注意的计算方法：

编码器：
在编码器中，将分辨率为H/4×W/4的c维标记化输入输入到连续的两个Swin Transformer块中进行表示学习，特征维度和分辨率保持不变。同时，patch 合并层会减少令牌的数量(2× down sampling)，将特征维数增加到2×原始维数。此过程将在编码器中重复三次。

patch合并层：
输入补丁被分为四部分，通过补丁合并层连接在一起。这样的处理会使特征分辨率下降2倍。并且，由于拼接操作的结果是特征维数增加了4倍，因此在拼接的特征上加一个线性层，将特征维数统一为原始维数的2倍。

Bottleneck：
由于Transformer太深，无法收敛，因此只使用连续两个Swin Transformer块构造瓶颈来学习深度特征表示。在瓶颈处，特征维度和分辨率保持不变。

解码器:
与编码器相对应的是基于Swin transformer模块的对称解码器。为此，与编码器中使用的patch 合并层不同，在解码器中使用patch扩展层对提取的深度特征进行上采样。patch 扩展层将相邻维度的特征图重塑为更高分辨率的特征图(2×上采样)，并相应地将特征维数减半。

补丁扩展层:
以第一个patch扩展层为例，在上采样之前，在输入特征上加一个线性层(W/32×H/32×8C)，将特征维数增加到原始维数的2倍(W/32×H/32×16C)。然后，利用重排操作将输入特征的分辨率扩展为输入分辨率的2倍，将特征维数降低为输入维数的1/4 (W/32×H/32×16C→W/16×H/16×4C)。

跳跃连接：
与U-Net类似，跳跃连接用于融合来自编码器的多尺度特征与上采样特征。将浅层特征和深层特征连接在一起，以减少降采样带来的空间信息损失。然后是一个线性层，连接特征的尺寸保持与上采样特征的尺寸相同。

实验结果：
数据集：
Synapse多器官分割数据集(Synapse):包括30例3779张腹部轴向临床CT图像。


不同方法对Synapse多器官CT数据集的分割结果

ACDC数据集上不同方法的分割精度

消融实验：
为了探讨不同因素对模型性能的影响，在Synapse数据集上进行了消融研究。具体来说，下面讨论了上采样、跳跃连接的数量、输入尺寸和模型尺度。

上采样的影响：
针对编码器中的patch 合并层，作者在解码器中专门设计了patch 扩展层，用于上采样和特征维数增加。为了探索所提出的补丁扩展层的有效性，在Synapse数据集上进行了双线性插值、转置卷积和补丁扩展层的Swin-Unet实验。下表中的实验结果表明，本文提出的Swin-Unet结合patch 扩展层可以获得更好的分割精度。

跳跃连接数目的影响：
Swin-UNet的跳跃连接被添加在1/4、1/8和1/16分辨率尺度的地方。通过将跳跃连接数分别更改为0、1、2和3，探讨了不同跳跃连接对模型分割性能的影响。从表4中可以看出，模型的分段性能随着跳过连接数的增加而提高。因此，为了使模型更加鲁棒，本工作中设置跳跃连接数为3。

输入大小的影响：
Swin-Unet在224 × 224、384 × 384输入分辨率下的测试结果如表5所示。随着输入尺寸从224 × 224增加到384 × 384，且patch尺寸保持4不变，Transformer的输入token序列会变大，从而提高模型的分割性能。然而，虽然模型的分割精度略有提高，但整个网络的计算负荷也有了显著增加。为了保证算法的运行效率，本文的实验以224 × 224的分辨率尺度作为输入。

模型大小的影响：
讨论了网络深化对模型绩效的影响。从表6可以看出，模型规模的增加并没有提高模型的性能，反而增加了整个网络的计算代价。考虑到精度和速度的权衡，采用基于tiny的模型进行医学图像分割。

讨论：
众所周知，模型前训练严重影响了基于transformer模型的性能。在本工作中，作者直接使用Swin-transformer在ImageNet上的训练权值来初始化网络编码器和解码器，这可能是一个次优方案。这种初始化方法是一种简单的方法，在未来作者表示将探索如何预处理Transformer端到端的医学图像分割。此外，由于本文输入的图像是二维的，而大部分医学图像数据是三维的，因此我们将在接下来的研究中探讨swin-unet在三维医学图像分割中的应用。

总结：
这篇文章介绍了一种新的基于纯transformer的u型编码解码器用于医学图像分割。标记化的图像补丁被输入到基于transformer的U形编码器-解码器结构中，带有跳跃连接，用于局部全局语义特征学习。具体来说，使用带有移位窗口的分层Swin transformer作为编码器来提取上下文特征。设计了一种基于Swin transformer的带有patch扩展层的对称解码器进行上采样操作，恢复特征图的空间分辨率。在对输入输出进行4×直接下采样和上采样的情况下，对多器官和心脏分割任务的实验表明，纯基于transformer的u形编解码网络的性能优于全卷积或transformer与卷积相结合的方法。为了充分利用Transformer的强大功能，作者将Swin Transformer块作为特征表示和远程语义信息交互学习的基本单元。在多器官和心脏分割任务上的实验表明，所提出的swin - unet具有良好的性能和泛化能力。