TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
发表时间：2021年2月
发表期刊：Arxiv

Abstract

医学图像分割是发展医疗保健系统的必要前提，尤其是疾病诊断和治疗计划。在各种医学图像分割任务上，u形架构，也称为U-Net，已经成为事实上的标准并取得了巨大的成功。然而，由于卷积操作的固有局部性，U-Net 通常在显式建模远程依赖方面表现出局限性。专为序列到序列预测而设计的 Transformer 已成为具有先天全局自我注意机制的替代架构，但由于低级细节不足，可能导致定位能力有限。在本文中，我们提出了 TransUNet，它兼具 Transformer 和 U-Net，作为医学图像分割的强大替代方案。一方面，Transformer 将来自卷积神经网络 (CNN) 特征映射的标记化图像块编码为用于提取全局上下文的输入序列。另一方面，解码器对编码特征进行上采样，然后将其与高分辨率 CNN 特征图相结合，以实现精确定位。

我们认为 Transformer 可以作为医学图像分割任务的强大编码器，结合 U-Net 通过恢复局部空间信息来增强更精细的细节。 TransUNet 在不同的医学应用（包括多器官分割和心脏分割）上实现了优于各种竞争方法的性能。代码和模型可在 https://github.com/Beckschen/TransUNet 获得。

Introduction

卷积神经网络（CNN），尤其是全卷积网络（FCN）[8]，已经在医学图像分割中占据主导地位。在不同的变体中，U-Net [12] 由具有跳跃连接的对称编码器-解码器网络组成，以增强细节保留，已成为事实上的选择。基于这种方法，在广泛的医学应用中取得了巨大的成功，例如磁共振 (MR) [16] 的心脏分割、计算机断层扫描 (CT) [7, 17, 19] 的器官分割和息肉从结肠镜检查视频中分割 [20]。

尽管基于CNN的方法具有非凡的表示能力，但由于卷积运算的固有局部性，它在显式长程关系建模方面普遍表现出局限性。因此，这些体系结构通常产生较弱的性能，特别是对于目标结构，显示在纹理，形状和大小方面患者间的巨大变化。为了克服这一局限性，现有的研究提出建立基于CNN特征的自我注意机制[13,15]。另一方面，为序列预测而设计的Transformer已经成为完全使用分配卷积算子并且仅仅依赖注意力机制的替代架构[14]。与先前的基于CNN的方法不同，Transformer不仅在建模全局上下文方面功能强大，而且在大规模预训练情况下对下游任务具有优越的可转移性。机器翻译和自然语言处理(nlp)领域已经广泛见证了这一成功[3,14]。最近，对于各种图像识别任务的尝试也达到甚至超过了最先进的表现水平[4,18]。

在本文中，我们提出了第一个研究，探讨了Transformer在医疗图像分割的背景下的潜力。然而，有趣的是，我们发现一种简单的用法(例如，使用Transformer对标记化的图像patch进行编码，然后直接将隐藏的特征表示上取样到一个高分辨率的密集输出中)并不能产生令人满意的结果。

这是因为Transformer将输入视为1d 序列，并且在所有阶段都专注于建模全局上下文，从而导致低分辨率的特征缺乏详细的定位信息。直接上采样不能有效地恢复到完全分辨率，从而导致粗分割结果。另一方面，CNN架构(例如，u-net [12])提供了一个提取低层次视觉线索的途径，可以很好地补救这些精细的空间细节。

为此，我们提出了 Transunet，第一个医学图像分割框架，它从序列到序列预测的角度建立了自我注意机制。为了弥补Transformer带来的特征分辨率的损失，Transunet 采用了一种混合的 CNN Transformer结构，利用CNN特征提供的详细的高分辨率空间信息和Transformer编码的全局上下文。受到 u 形结构设计的启发，Transformer编码的自我关注特征被上采样，然后与跳过编码路径的不同高分辨率CNN特征相结合，以实现精确定位。我们证明这样的设计使我们的框架能够保留Transformer的优点，也有利于医疗图像分割。实验结果表明，与以前的基于CNN的自我注意方法相比，我们的基于Transformer的体系结构提供了一种更好的利用自我注意的方法。此外，我们观察到更密集的低水平特征结合通常导致更好的分割准确性。广泛的实验证明了我们的方法在各种医疗图像分割任务中相对于其他竞争方法的优越性。

Related Works

Combining CNNs with self-attention mechanisms
各种研究试图通过建立基于特征映射的全局像素相互作用模型，将自我注意机制整合到卷积神经网络中。例如，Wang 等人设计了一个非局部算子，它可以插入多个中间卷积层[15]。Schlemper 等[13]建立在编解码器 u 形结构的基础上，提出了集成到跳跃连接中的加性注意门模块。与这些方法不同的是，我们使用Transformer将整体自我注意嵌入到我们的方法中。
Transformers
Transformer最初由[14]提出用于机器翻译，并在许多 nlp 任务中建立了最新的技术状态。为了使Transformer也适用于计算机视觉任务，已经进行了几次修改。例如，parmar 等[11]只在每个查询像素的局部邻域应用自我注意，而不是全局应用。Child 等[1]提出了稀疏Transformer，它采用可扩展的近似来实现全局自我注意。最近，Vision Transformer(vit)[4]通过将全局自注意Transformer直接应用于全尺寸图像，在图像网分类方面取得了最新进展。据我们所知，建议的 transunet 是第一个基于Transformer的医疗图像分割框架，它建立在非常成功的 vit 的基础上。

Method

给出一幅图像 x ∈ R^h × w × c^，其空间分辨率为 h × w，通道数为 c。我们的目标是预测相应的大小为 h × w 的像素级标签图。最常见的方法是直接训练CNN(例如 unet)来首先将图像编码成高级特征表示，然后将其解码回完整的空间分辨率。与现有的方法不同，我们的方法通过使用Transformer将自注意机制引入到编码器设计中。我们将在3.1节中首先介绍如何直接应用Transformer来对分解后的图像patch中的特征表示进行编码。然后，在第3.2节中将详细阐述 Transunet 的总体框架。

Transformer as Encoder

Image Sequentialization
在文献[4]之后，我们首先通过将输入 x 重新形成一个展平的2d patch序列{ xⁱ p ∈ R^p2·c | i = 1……n，来进行标记化} ，其中每个patch的大小为 p × p，n = hw/p²是图像patch的数量(即输入序列长度)。
Patch Embedding
使用可训练的线性投影将 xp 矢量化为潜在的 d 维嵌入空间。为了对patch空间信息进行编码，我们学习了添加到patch嵌入中以保留位置信息的特定位置嵌入，如下所示:

其中，E∈R^(p2 · C)×D^ 是patch嵌入投影，Epos ∈ R^N×D表示位置嵌入。

Transformer 编码器由多头自注意(MSA)和多层感知器(MLP)块(Eq.(2)(3).因此l-th 层的输出可以写成如下:

其中，LN（）表示层归一化运算符，zl 是编码后的图像表示。Transformer层的结构如图1(a)所示。

Fig.1 (a)Transformer层的示意图; (b)建议的 Transunet.patch 的结构

TransUnet

为了达到预测密集输出的目的，一个直观的解决方案是简单地将编码特征表示 zl ∈ R^hw/p2×D提高到全分辨率。在这里，为了恢复空间序列，编码特征的大小应该首先由 hw/p2重塑为 hp × wp。我们使用1 × 1卷积将重构特征的通道大小降低为类数，然后将特征映射直接双线性上采样至全分辨率 h × w，以预测最终的分割结果。在4.3节后面的比较中，我们将这个初始的上采样基线表示为解码器设计中的“无”。

尽管将Transformer与初始上采样相结合已经产生了合理的性能，如上所述，但这种策略并不是变压器在分割中的最佳使用，因为 hp × wp 通常比原始图像分辨率 h×w 小得多，因此不可避免地导致低层细节(如器官的形状和边界)的丢失。因此，为了补偿这种信息损失，Transunet 采用了一种混合的CNN Transformer结构作为编码器以及级联的上采样器来实现精确定位。所提议的Transformer的概述如图1所示。
CNN-Transformer Hybrid as Encoder
Transunet 没有使用纯粹的Transformer作为编码器(3.1节) ，而是使用了 cnn-transformer 混合模型，其中 cnn 首先被用作特征提取器，为输入生成一个特征映射。对从 cnn 特征映射中提取的图像进行了1×1 patch嵌入，取代了从原始图像中提取的patch嵌入方法。

我们选择这种设计，因为1)它允许我们在解码路径中利用中高分辨率 cnn 特征映射; 2)我们发现混合 cnn-transformer 编码器比简单地使用纯transformer作为编码器性能更好。
Cascaded Upsampler
我们提出了一种级联上采样器(CUP) ，它由多个上采样步骤组成，对隐藏特征进行解码，输出最终的分割掩模。将隐藏特征 zl ∈R^HW/P2 × D^ 的序列重构成 HP × WP × D 的形状后，通过级联多个上采样块实例化CUP，实现从 HP×WP 到 HxW 的完全分辨率，其中每个块由2 × 上采样算子、3 × 3卷积层和 Relu 层依次组成。

我们可以看到CUP与混合编码器形成一个 u 形架构，通过跳过连接，能够在不同分辨率水平上聚合特性。CUP的详细结构以及中间的跳跃式连接可以在图1(b)中找到。

Experiments and Discussion

Dataset and Evaluation

突触多器官分割数据库1。我们在 miccai 2015多层螺旋 ct 腹部标记挑战中使用了30次腹部 ct 扫描，共计3779次轴向增强腹部临床 ct 图像。

每个CT体积由85～198个512×512像素的切片组成，体素空间分辨率为([0.54～0.54]×[0.98～0.98]×[2.5～5.0])mm3。 在[5]的基础上，我们报告了18例（2212个轴位切片）和12例验证的腹部8个器官（主动脉、胆囊、脾脏、左肾、右肾、肝、胰、脾、胃）的平均DSC和平均豪斯距离(HD)。
自动心脏诊断挑战2。 ACDC挑战收集从MRI扫描仪获得的DI患者的检查。 在屏气状态下获得电影MR图像，一系列短轴切片覆盖心脏，从左心室底部到心尖，层厚5～8毫米。 短轴面内空间分辨率为0.83～1.75mm2/像素。

每个病人的扫描都用左心室(LV)、右心室(RV)和心肌(MYO)的ground truth手工注释。 我们报告平均DSC随机拆分70个训练案例（1930个轴向切片），10个案例进行验证，20个案例进行测试。

Implementation Details

对于所有的实验，我们应用简单的数据扩充，例如，随机旋转和翻转。

对于纯基于Transformer的编码器，我们简单地采用了12个Transformer层的VIT[4]。 对于混合编码器的设计，本文将RESNET-50[6]和VIT相结合，称为“R50-VIT”。 所有Transformer主干（即VIT）和Resnet-50（表示为“R-50”）在ImageNet上预先训练[2]。 输入分辨率和patch大小P设置为224×224和16，除非另有说明。 因此，我们需要在CUP中连续级联四个2×上采样块以达到全分辨率。 并用SGD优化器对模型进行训练，其学习率为0.01，动量为0.9，重量衰减为1e-4。 ACDC数据集和Synapse数据集的训练迭代次数分别为20K和14K。 所有的实验都是使用一个单一的NVIDIA RTX2080TI GPU进行的。

在[17，19]之后，以逐个切片的方式推断所有3D体积，并且将预测的2D切片堆叠在一起以重建用于评估的3D预测。

Comparison with State-of-the-arts

我们在突触多器官分割数据集上进行了主要的实验，将我们的Transunet与前四个最先进的数据集进行了比较:1)V-Net[9]； 2)DARR[5]； 3)U-Net[12]和4)Attnunet[13]。

为了验证CUP译码器的有效性，我们使用VIT[4]作为编码器，并分别比较了朴素上采样（“none”）和CUP译码器的结果； 为了验证我们设计的混合编码器的有效性，我们使用CUP作为解码器，并分别比较了VIT和R50VIT作为编码器的结果。 为了使与ViT混合基线(R50-Vit-Cup)和我们的Transunet进行比较的更加公平，我们还用ImageNet预先训练的Resnet-50替换了U-Net[12]和Attnunet[10]的原始编码器。 DSC和平均Hausdor距离（以毫米为单位）的结果在表1中报告。

首先，与VIT-NONE相比，VIT-CUP的平均DSC和Hausdor距离分别提高了6.36%和3.50mm。 这一改进表明我们的CUP设计提供了比直接上采样更好的解码策略。 与VIT-CUP相比，R50-VIT-CUP的DSC性能提高了3.43%,豪斯距离提高了3.24mm,证明了我们的混合编码器的性能。 我们的Transunet建立在R50-VIT-Cup的基础上，它还配备了跳过连接，在基于Transformer的模型中实现了最好的结果。

其次，表1还表明，与现有技术相比，所提出的Transunet有显著的改进，例如，考虑到平均DSC，性能增益从1.91%到8.67%不等。 特别地，直接应用Transformer进行多器官分割的结果是合理的（VIT-Cup的DSC为67.86%)，但不能与U-NET或Attnunet相媲美。 这是由于Transformer能很好地捕捉高层语义，这有利于分类任务，但缺乏低层线索，以分割医学图像的精细形状。 另一方面，将Transformer与CNN相结合，即R50-VitCup的性能优于V-NET和DARR，但仍低于单纯基于CNN的R50-U-NET和R50-Attnunet。 最后，将Transunet通过跳过连接与U-NET结构相结合，使其达到了新的水平，比R50-Vit-Cup和R50-Attnunet分别提高了6.19%和1.91%，显示了Transunet在医学图像分割中学习高层语义特征和底层细节的强大能力，这在医学图像分割中是至关重要的。 平均Hausdor距离也有类似的趋势，这进一步证明了我们的Transunet优于这些基于CNN的方法。

Analytical Study

为了彻底评估所提出的Transunet框架并验证其在特定环境下的性能，进行了多种消融研究，7包括：1）跳过连接的数量； 2）输入分辨率； 3）序列长度和patch大小；4）模型缩放。

The Number of Skip-connections
如上所述，通过恢复低级空间信息，集成类似U-Net跳过连接有助于增强更精细的分割细节。 这次消融实验的目的是测试在Transunet中增加跳过连接数的影响。 通过改变跳过连接的数量为0(R50-Vit-Cup)/1/3，所有8个测试器官的平均DSC分割性能总结在图2中。 注意，在“1-skip”设置中，我们只在1/4分辨率范围内添加跳过连接。 我们可以看到，添加更多的跳过连接通常会导致更好的分割性能。 最佳的平均DSC和HD是通过插入跳过连接到CUP的所有三个中间上采样步骤，除了输出层，即在1/2、1/4和1/8分辨率尺度（如图1所示）。 因此，我们对Transunet采用这种配置。 同样值得一提的是，较小器官（即主动脉、胆囊、肾脏、胰腺）的性能增益比较大器官（即肝、脾、胃）更为明显。 这些结果加强了我们最初的直觉，即在Transformer设计中集成类似U-网的跳过连接，从而能够学习精确的低级细节。

作为一个有趣的研究，我们将加法变换应用到跳过连接中，类似于[13]，并发现这种新型的跳过连接可以进一步提高分割性能。 由于GPU内存的限制，我们在1/8分辨率尺度的跳过连接中使用了一个轻量级transformer，而保持其他两个跳过连接不变。 结果，这个简单的改变导致性能提高1.4%的DSC。

Fig.2 Transunet中跳过连接数的消融研究。

On the Influence of Input Resolution
Transunet的默认输入分辨率为224×224。 这里，我们还提供了在高分辨率512×512上训练Transunet的结果，如表2所示。 当采用512×512作为输入时，我们保持相同的patch大小（即16)，使得Transformer的平均序列长度近似为5×Transformer序列长度。 正如[4]所指出的，增加有效序列长度显示出鲁棒性的改进。 对于Transunet，将分辨率尺度从224×224改变到512×512，平均DSC提高了6.88%，但代价是计算量要大得多。 因此，考虑到计算成本，本文的所有实验比较都是在默认分辨率为224×224的情况下进行的，以证明Transunet的有效性。

On the Influence of Patch Size/Sequence Length
我们还考察了patch大小对Transunet的影响。 结果汇总在表3中。 研究表明，较小的patch大小通常可以获得较高的分割性能。 注意，transformer的序列长度与patch大小的平方成反比（例如，patch大小16对应于序列长度196，而patch大小32具有较短的序列长度49)，因此减小patch大小（或增加有效序列长度）显示出稳健的改进，因为Transformer编码用于较长输入序列的每个元素之间更复杂的依赖关系。 遵循VIT[4]中的设置，我们在本文中使用16×16作为默认补丁大小。

Model Scaling
最后，我们提供了关于Transunet模型尺寸的消融研究。 特别地，我们研究了两个直接的Transunet构型，“基”和“大”模型。 对于“基本”模型，隐藏尺寸D、层数、MLP尺寸和头数分别设置为12、768、3072和12，而对于“大型”模型，超参数分别为24、1024、4096和16。 从表4我们得出结论，更大的模型会导致更好的性能。 考虑到计算量的问题，我们对所有的实验都采用了“基础”模型。

Visualization

我们提供了Synapse数据集中的定性比较结果，如图3所示。 可以看出：1）纯基于CNN的方法U-Net和Attnunet更容易对器官进行过分割或欠分割（例如，在第二行中，脾脏被Attnunet过度分割，而被Unet欠分割），这表明基于Transformer的模型，如我们的Transsunet或R50-Vit-Cup具有更强的全局上下文编码和语义区分能力。 2）第一行的结果表明，我们的Transunet比其他方法预测的假阳性更少，这表明Transunet在抑制这些噪声预测方面比其他方法更有优势。 3）为了在基于Transformer的模型中进行比较，我们可以观察到R50-VIT-Cup的预测在边界和形状方面比Transunet的预测更粗糙（例如第二行胰腺的预测）。 此外，在第三行，Transunet正确地预测了左肾和右肾，而R50-Vit-Cup错误地填充了左肾的内孔。 这些观察表明Transunet能够更好地分割和保存详细的形状信息。 原因是Transunet既有高级全局上下文信息，又有低级细节，而R50-Vit-Cup只依赖于高级语义特性。 这再次验证了我们最初的直觉，将U型网络般的跳过连接集成到变压器设计中，以实现精确的定位。

Fig.3 可视化方法定性比较。 从左到右：(a)ground truth，(b)Transunet，©R50-Vit-Cup，(d)R50attnunet，(e)R50-U-net。 我们的方法预测较少的假阳性，保持较好的信息。

Generalization to Other Datasets

为了显示Transunet的泛化能力，我们进一步评估了其他成像模式，即旨在自动心脏分割的MR数据集ACDC。 我们观察到Transunet相对于单纯基于CNN的方法(R50-UNET和R50-ATTNUNET)和其他基于Transformer的基线（VIT-CUP和R50-VIT-CUP）的持续改进，这与以前在Synapse CT数据集上的结果相似。

Conclusion

Transformer被称为具有强大的先天自我注意机制的架构。 在这篇论文中，我们提出了第一个研究，以调查Transformer在一般医学图像分割中的使用。 为了充分利用Transformers的强大功能，Transunet不仅通过将图像特征作为序列来编码强全局上下文，而且通过U型混合结构设计很好地利用了低层CNN特征。 Transunet作为一种替代基于FCN的医学图像分割方法的框架，其性能优于各种竞争方法，包括基于CNN的自关注方法。