图像分割（三）—— Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

Abstract

在过去的几年中，卷积神经网络(CNNs)在医学图像分析中取得了里程碑。特别是基于u型结构和跳跃连接的深度神经网络在各种医学图像任务中得到了广泛的应用。然而，CNN虽然取得了优异的性能，但由于卷积操作的局部性，它并不能很好地学习全局和长期的语义信息交互。在本文中，我们提出了Swin-Unet，这是一个类似于unet的纯变压器的医学图像分割。将标记化的图像补丁输入到基于变压器的u形编码器-解码器架构中，用于局部全局语义特征学习。具体来说，我们使用带有移位窗口的分层摇摆变换器作为编码器来提取上下文特征。设计了一种基于补丁扩展层的对称双变换解码器，进行上采样操作，恢复特征图的空间分辨率。在4×的输入和输出直接下采样和上采样下，多器官和心脏分割实验证明了这一点。

Introduction

得益于深度学习的发展，计算机视觉技术在医学图像分析中得到了广泛的应用。图像分割是医学图像分析的重要组成部分。特别是，准确和稳健的医学图像分割可以在计算机辅助诊断和图像引导的临床手术中发挥基石作用[1，2]。

现有的医学图像分割方法主要依赖于U型结构的全卷积神经网络[3，4，5]。典型的U型网络，U-Net[3]，由一个带有跳跃连接的对称编解码器组成。在编码器中，采用一系列卷积层和连续下采样层来提取大感受野的深度特征。然后，解码器将提取的深度特征上采样到输入分辨率进行像素级语义预测，并将来自编码器的不同尺度的高分辨率特征通过跳跃连接进行融合，以缓解下采样造成的空间信息损失。凭借这种优雅的结构设计，U-Net在各种医学成像应用中取得了巨大的成功。根据这一技术路线，已经开发了许多算法，如3DU-NET[6]、RES-UNET[7]、U-NET++[8]和UNET3+[9]，用于各种医学成像模式的图像和体积分割。这些基于FCNN的方法在心脏分割、器官分割和病变分割中都表现出很好的性能，证明了CNN具有很强的区分特征学习能力。

目前，基于CNN的分割方法虽然在医学图像分割领域取得了很好的效果，但仍不能完全满足医学应用对分割精度的严格要求。图像分割仍然是医学图像分析中的一项具有挑战性的任务。由于卷积运算的固有局部性，基于CNN的方法很难学习显式的全局和远程语义信息交互[2]。一些研究试图通过使用Arous卷积层[10，11]、自我注意机制[12，13]和图像金字塔[14]来解决这个问题。然而，这些方法在对长期依赖关系进行建模方面仍然存在局限性。最近，受Transformer在自然语言处理(NLP)领域的巨大成功的启发[15]，研究人员试图将Transformer带入视觉领域[16]。在[17]中，提出了视觉转换器(ViT)来执行图像识别任务。以位置嵌入的二维图像块为输入，在大数据集上进行预训练，VIT取得了与基于CNN的方法相当的性能。此外，文献[18]中提出了数据高效的图像转换器(DeiT)，这表明Transformer可以在中等大小的数据集上进行训练，并与蒸馏方法相结合可以得到更健壮的Transformer。在[19]中，开发了一种分层Swin Transformer。文[19]以Swin Transformer为视觉骨干，在图像分类、目标检测和语义分割等方面取得了一流的性能。ViT、DeiT和Swin Transformer在图像识别任务中的成功展示了Transformer在视觉领域的应用潜力。

受到Swin Transformer[19]成功的激励，我们建议Swin-UNET在这项工作中利用Transformer的力量进行2D医学图像分割。据我们所知，Swin-UNET是第一个纯基于transformer的U型架构，由编码器、bottleneck、解码器和跳跃连接组成。编码器、bottleneck和解码器都是基于Swin Transformer块[19]构建的。输入的医学图像被分割成不重叠的图像块。每个patch都被视为一个令牌，并被馈送到基于Transformer的编码器中，以学习深层的特征表示。提取的上下文特征由带patch expanding layer的解码器进行上采样，并通过跳跃连接与来自编码器的多尺度特征进行融合，从而恢复特征地图的空间分辨率并进一步进行分割预测。在多器官和心脏分割数据集上的大量实验表明，该方法具有良好的分割精度和较强的泛化能力。

我们的贡献可以概括为：

(1)基于Swin Transformer块，构建了一种跳跃连接的对称编解码器结构。在编码器中，实现从局部到全局的自关注；在解码器中，将全局特征向上采样到输入分辨率，以进行相应的像素级分割预测。

(2)在不使用卷积和插值运算的情况下，开发了一种patch expanding layer，实现了上采样和特征维数的增加。

(3)实验中发现跳跃连接对transformer同样有效，因此最终构建了一种基于transformer的纯U型跳跃连接编解码器体系结构Swin-UNET(Swin-UNET)(简称Swin-UNET)。

Method

3.1 Architecture overview

提出的的Swin-UNET的总体架构如图1所示。Swin-UNET由编码器、bottleneck、解码器和跳跃连接组成。Swin-UNET的基本单元是Swin Transformer block[19]。对于编码器，为了将输入转化为序列嵌入，将医学图像分割成大小为4×4的互不重叠的块，通过这种划分方法，每个块的特征维数变为4×4×3=48. 此外，将线性嵌入层应用于将特征维度投影到任意维度(表示为C)。转换后的patch令牌通过几个Swin Transformer block和patch merging layer来生成分层特征表示。具体地说，patch merging layer负责下采样和增维，Swin Transformer block负责特征表示学习。受U-Net[3]的启发，我们设计了一种基于对称变压器的解码器。解码器由Swin Transformer block 和patch expanding layer组成。提取的上下文特征通过跳跃连接与来自编码器的多尺度特征融合，以弥补下采样造成的空间信息损失。与patch merging layer不同的是，专门设计了patch expanding layer来执行上采样。patch expanding layer将相邻维度的特征图重塑为分辨率为2倍上采样的大型特征图。最后，利用最后一层patch expanding layer进行4倍上采样，将特征图的分辨率恢复到输入分辨率(W×H)，然后在这些上采样的特征上应用线性投影层输出像素级分割预测。

3.2 Swin Transformer block

与传统的多头自关注(MSA)模块不同，Swin-Transformer block[19]是基于移位窗口构造的。在图2中,提出了两个连续的swin transformer block。每个swin transformer block由LN层、多头自关注模块、残差连接和具有Gelu非线性的二层MLP组成。基于窗口的多头自关注(W-MSA)模块和基于移位窗口的多头自关注(SW-MSA)模块分别应用于连续的两个变压器块。基于这种窗口划分机制，连续的swin transformer block可以表示为：

Encoder
在编码器中，将分辨率为H/4×W/4的C维标记化输入送入两个连续的Swin Transform block进行表示学习，其中特征尺寸和分辨率保持不变。同时，patch expanding layer将减少表征数(2倍下采样)，并将特征维数提高到原来的2倍。此过程将在编码器中重复三次(dimension和resolution的变化是在patch merging layer完成的，2个Swin Transformer block+patch merging layer是一组，一共重复3次)。

Patch merging layer
输入的patch被分成4个部分，并通过patch merging layer连接在一起。经过这样的处理，特征分辨率将降低2倍。并且，由于拼接操作导致特征维数增加了4倍，因此在拼接的特征上应用线性层以将特征维数统一到2倍的原始维数(patch merging layer后（或中）有一个linear layer，作用：将dimension的4改为2)。

Bottleneck
由于transformer太深而无法收敛[33]，因此只使用两个连续的Swin Transformer block来构建学习深层特征表示的bottleneck。在bottleneck中，特征尺寸和分辨率保持不变。

Decoder
与编码器相对应，对称解码器是基于Swin Transformer块构建的。为此，与编码器中使用的patch merging layer不同，我们在解码器中使用patch expanding layer对提取的深度特征进行上采样。patch expanding layer将相邻维度的特征图重塑为更高分辨率的特征图(2倍上采样)，并相应地将特征维数降至原始维数的一半。

Patch expanding layer
以第一个patch expanding layer为例，在上采样之前，对输入特征(W /32×H/32×8C)应用linear layer，将特征维数增加到原来的2倍(W/32×H/32×16C)。然后，使用rearrange operation将输入特征的分辨率扩展到输入分辨率的2倍，并将特征维数降低到输入维数的四分之一(W/32×H/32×16C→W/16×H/16×4C)。我们将在4.5节讨论使用patch expanding layer执行上采样的影响。

Skip connection

与U-Net[3]类似，跳跃连接用于将编码器中的多尺度特征与上采样特征融合。我们将浅层特征和深层特征拼接在一起，以减少下采样造成的空间信息损失。在linear layer之后，串联特征的维度保持与上采样特征的维度相同。在第4.5节中，我们将详细讨论跳跃连接的数量对模型性能的影响