Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

[深度学习论文笔记]Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation
Swin-Unet：用于医学图像分割的类Unet纯transformer
论文：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的编码器中，以学习深度特征表示。提取的上下文特征由带补丁扩展层的解码器上采样，通过跳跃连接与编码器的多尺度特征融合，恢复特征图的空间分辨率，进一步进行分割预测。在多器官和心脏数据集上的大量实验表明，该方法具有良好的分割精度和鲁棒泛化能力。

总结：
  卷积神经网络（CNN）在医学图像分析中取得了里程碑式的进展。尤其是，基于U形架构和跳跃连接的深度神经网络已广泛应用于各种医学图像任务中。但是，尽管CNN取得了出色的性能，但是由于卷积操作的局限性，它无法很好地学习全局和远程语义信息交互。

  作者提出了Swin-Unet，它是用于医学图像分割的类似Unet的纯Transformer。标记化的图像块通过跳跃连接被馈送到基于Transformer的U形En-Decoder架构中，以进行局部全局语义特征学习。

具体来说，作者的贡献可以总结为:
  (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)，然后在这些上采样特征上应用线性投影层输出像素级分割预测。

总结：
  patch大小为4x4，因此每个patch的特征维度是4x4x3=48，之后应用线性嵌入层投影特征维度到任意维度（表示为C)。
  转换的patch token经过几个swin transformer块和patch合并层，以产生层级特征表示。patch合并层用于下采样和升维，swin transformer块用于特征表示学习。
受unet启发，设计对称的基于transformer的解码器，由swin transformer块和patch扩展层组成。提取的上下文特征通过跳跃连接于编码器的多尺度特征融合，以补偿下采样造成的空间信息损失。
  相比于patch合并层，一个patch扩展层变形毗邻维度的特征图到一个2倍上采样的大特征图。最后的patch扩展层用于直径4倍的上采样，以恢复特征图的分辨率到输入大小（wxh),然后线性投影层用于这些上采样的特征，输出像素级别的分割预测。

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

总结：
  具体来说，作者使用带有偏移窗口的分层Swin Transformer作为编码器来提取上下文特征。并设计了具有补丁扩展层的基于对称Swin Transformer的解码器来执行上采样操作，以恢复特征图的空间分辨率。

结构详解：
  两个连续的Swin Transformer块：

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

自注意的计算方法：

M2代表patch数量，d代表q或k的维度。B中的值来源于偏置矩阵。

编码器：
   C维的分辨率维H/4 x W/4的标记化输入输入到两个连续的Swin Transformer块中进行表示学习，特征维度和分辨率不变。同时，patch合并层将缩减token的数量（2倍下采样），提升特征维度到原始维度的2倍。这个步骤在编码器中重复3次。

patch合并层：
  输入patch分成4部分，通过patch合并层拼接在一起，采用这个过程，特征分辨率将下采样2倍，因为拼接操作导致特征维度提升4倍，一个线性层用于拼接的特征，将特征维数统一为2×原始维数

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

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

  基于swin transformer块构建，在解码器中采用patch扩展层，用于删改杨提取的特征，patch扩展层将相邻维度的特征图变为更高分辨率（2倍上采样），维度数量减半。

patch 扩展层：
  以第一个patch扩展层为例，上采样前，一个线性层应用到输入特征（W/32 x H/32 X 8C)以增加特征维度到原来的2倍（W/32 x H/32 X 16C)。然后，采用重排操作，扩展原输入特征的分辨率2倍，同时缩减特征维度到输入维度的1/4（W/32×H/32×16C->W/16×H/16×4C）

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

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


  不同方法对Synapse多器官CT数据集的分割结果（基于CNN的方法有过度分割的问题，这可能是卷积的局部性引起的）

  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具有良好的性能和泛化能力。