TransUNet：变压器成为用于医学图像分割的强大编码器
TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
作者单位：JHU, 电子科大, 斯坦福大学等
代码：Beckschen/TransUNet
论文：https://arxiv.org/abs/2102.04306
摘要
医学图像分割是开发医疗保健系统（尤其是疾病诊断和治疗计划）的必要先决条件。在各种医学图像分割任务中，u型架构（也称为U-Net）已成为事实上的标准，并取得了巨大的成功。但是，由于卷积运算的固有局部性，U-Net通常在明确建模远程依赖关系方面显示出局限性。设计用于序列到序列预测的变压器已经成为具有先天性全局自注意机制的替代体系结构，但由于底层细节不足，可能导致定位能力受到限制。在本文中，我们提出了TransUNet，它同时具有Transformers和U-Net的优点，是医学图像分割的强大替代方案。
- 一方面，Transformer将来自卷积神经网络（CNN）特征图的标记化图像块编码为提取全局上下文的输入序列。
- 另一方面，解码器对编码的特征进行上采样，然后将其与高分辨率的CNN特征图组合以实现精确的定位。
我们认为，借助U-Net的组合，通过恢复本地化的空间信息，可以将Transformers用作医学图像分割任务的强大编码器。 TransUNet在各种医疗应用（包括多器官分割和心脏分割）上均比各种竞争方法具有更高的性能。 代码和模型可在https://github.com/Beckschen/ TransUNet上获得。
1 引言
卷积神经网络（CNN），尤其是全卷积网络（FCN）[8]，已在医学图像分割中占主导地位。 在不同的变体中，U-Net [12]由具有跳跃连接的对称编码器-解码器网络组成，以增强细节保留，已成为事实上的选择。 基于这种方法，已在广泛的医学应用中取得了巨大的成功，例如磁共振（MR）的心脏分割[16]，计算机断层扫描（CT）的器官分割[7,17,19]和息肉 结肠镜检查视频分割[20]。（加入Unet文献综述）
尽管基于卷积神经网络的方法具有出色的表示能力，但由于卷积运算的固有局限性，通常在建模显式远程关系方面仍存在局限性。因此，这些体系结构通常产生较弱的性能，特别是对于目标结构，这些结构在纹理，形状和大小方面表现出较大的患者间差异。为了克服这一局限性，现有研究提议基于CNN特征建立自我注意机制[13,15]。另一方面，为序列到序列预测而设计的变压器已经出现，作为一种替代体系结构，它完全采用了分配卷积运算符，而仅依赖于注意力机制[14]。与以前的基于CNN的方法不同，Transformers不仅在建模全局上下文方面功能强大，而且在大规模的预训练下对下游任务也显示出卓越的可传递性。在机器翻译和自然语言处理（NLP）领域中已广泛见证了这一成功[3,14]。最近，针对各种图像识别任务的尝试也达到甚至超过了最新技术水平[4,18]。
在本文中，我们提出了第一个研究，在医学图像分割的背景下探讨了变压器的潜力。 但是，有趣的是，我们发现单纯的用法（即使用转换器对标记化的图像斑块进行编码，然后直接将隐藏的特征表示升采样为完整分辨率的密集输出）无法产生令人满意的结果。
这是由于Transformers将输入视为1D序列，并且仅专注于在所有阶段建模全局上下文，因此会导致缺乏详细的本地化信息的低分辨率功能。 而且，无法通过直接上采样到完整分辨率来有效地恢复此信息，因此会导致粗略的分割结果。 另一方面，CNN架构（例如，U-Net [12]）提供了提取低级视觉提示的途径，可以很好地弥补此类精细的空间细节。
为此，我们提出了TransUNet，这是第一个医学图像分割框架，它从序列到序列的预测角度建立了自我注意机制。为了补偿Transformers带来的特征分辨率的损失，TransUNet采用了混合CNN变形器体系结构，以利用来自CNN特征的详细高分辨率空间信息以及Transformers编码的全局上下文。受u形建筑设计的启发，然后对Transformers编码的自注意特征进行上采样，以与从编码路径中跳过的不同高分辨率CNN特征结合，以实现精确的定位。我们表明，这种设计使我们的框架能够保留Transformers的优势，也能有益于医学图像分割。实证结果表明，与以前的基于CNN的自我注意方法相比，基于Transformer的体系结构提供了一种更好的利用自我注意的方法。此外，我们观察到，更深入地合并低级特征通常会导致更好的分割精度。大量的实验证明了我们的方法在各种医学图像分割任务上相对于其他竞争方法的优越性。
2 相关研究
将CNN与自我注意机制结合起来。 各种研究已尝试通过基于特征图对所有像素的全局交互进行建模来将自我注意机制集成到CNN中。 例如，Wang等。 设计了一个非本地运算符，可以将其插入多个中间卷积层中[15]。 Schlemper等人建立在编解码器U形架构的基础上。 [13]提出了将附加注意门模块集成到跳过连接中。 与这些方法不同，我们使用“Transformers”将全局自我注意嵌入到我们的方法中。
Transformers.。 [14]首先提出了用于机器翻译的变压器，并在许多NLP任务中建立了最先进的技术。 为了使“Transformers”也适用于计算机视觉任务，已进行了一些修改。 例如，Parmar等。 [11]仅针对每个查询像素在局部邻域而不是全局中应用了自我注意。 儿童等。 [1]提出了稀疏变压器，该变压器采用可缩放的近似来实现全局自注意。 最近，Vision Transformer（ViT）[4]通过将具有全局自我关注能力的Transformers直接应用于全尺寸图像，从而实现了ImageNet分类的最新技术。 据我们所知，所提议的TransUNet是第一个基于Transformer的医学图像分割框架，该框架基于非常成功的ViT建立。
3 方法
给定图像x∈RH×W×C，其空间分辨率为H×W，通道数为C。 我们的目标是预测尺寸为H×W的相应像素级标签图。最常见的方法是直接训练CNN（例如UNet），首先将图像编码为高级特征表示，然后将其解码回 完整的空间分辨率。 与现有方法不同，我们的方法通过使用变压器将自我关注机制引入编码器设计。 我们将首先在第3.1节中介绍如何直接应用转换器对来自分解后的图像块的特征表示进行编码。 然后，将在3.2节中详细说明TransUNet的总体框架。
3.1 变压器
***作为编码器的图像序列化***在[4]之后，我们首先通过将输入x重塑为扁平的2D补丁序列来执行标记化(轴∈RP2·C|和=1，…，N}，其中每个补丁的大小是P×P，N=硬件P2是图像补丁的数量（即输入序列长度），其中，每个补丁的尺寸是图像块的数目（即输入序列长度） P × P P \times P P×P和 N = H W / P − 2 N=HW/P-{2} N=HW/P−2
图 1：框架概述。 (a)变压器层示意图；(b)拟议的跨网结构。
Patch Embedding. 我们使用可训练的线性投影将矢量化斑块XP映射到一个潜在的D维嵌入空间。 为了对补丁空间信息进行编码，我们学习了添加到补丁嵌入中的特定位置嵌入，以保留位置信息如下：
变压器编码器由L层多头自注意(MSA)和多层感知器(MLP)块（等式）组成（2)(3)）。因此，`层的输出可写入如下：
其中LN（·）表示层归一化运算符， z l z_{l} zl为编码图像表示。变压器层的结构如图1(a)所示。
为了分割的目的，一个直观的解决方案是简单地将编码的特征表示 z L ∈ H W / P 2 × D z_{L}∈HW/P^{2}×D zL∈HW/P2×D升采样到全分辨率，以预测密集输出。 在这里，为了恢复空间顺序，编码特征的大小应该首先从 H W / P 2 HW/P^{2} HW/P2重塑到 H / P × W / P H/P×W/P H/P×W/P。为了将重塑特征的信道大小减少为类数，然后直接将特征映射上采样到全分辨率的 H × W H×W H×W，以预测最终的分割结果。在后面第4.3节的比较中，我们在解码器设计中将这个朴素的上采样基线表示为“无”。
虽然将变压器与朴素上采样相结合已经产生了合理的性能，但如上所述，这种策略并不是变压器在分割中的最佳使用，因为 H / P × W / P H/P×W/P H/P×W/P通常比原始图像分辨率 H × W H×W H×W小得多，因此不可避免地导致低级别细节的丢失（例如器官的形状和边界）。 因此，为了补偿这种信息损失，TransUnet采用了一种混合的CNN-Transformer体系结构作为编码器以及级联的upsampler，以实现精确的定位。 拟议的跨UNET概述如图1所示.
CNN-Transformer Hybrid as Encoder。 而不是使用纯变压器作为编码器（3.1节），TransUnet采用CNN-Transformer混合模型，其中CNN首先用作特征提取器，为输入生成特征映射。 贴片嵌入应用于从CNN特征映射中提取的1×1个补丁，而不是从原始图像中提取。 我们选择这种设计，因为
1)它允许我们利用解码路径中的中间高分辨率CNN特征映射；
2)我们发现混合CNN-Transformer编码器的性能优于简单地使用纯变压器作为编码器。
Cascaded Upsampler。我们引入了一个级联的上采样器，它包括多个上采样步骤来解码输出最终分割掩码的隐藏特征。重塑隐藏特性的序列 z L ∈ H W / P 2 × D z_{L}∈HW/P_{2}×D zL∈HW/P2×D，我们实例化杯通过级联多个上采样块达到全分辨率从 H / P × W / P H/P×W/P H/P×W/P到 H × W H×W H×W，其中每个块包括2×上采样操作符，3×3卷积层，和ReLU层连续。
我们可以看到，CUP与混合编码器一起形成了一个u形体系结构，它通过跳接连接在不同的分辨率级别上实现特征聚合。 CUP的详细架构以及中间跳过连接可以在图1(b)中找到)。
4实验和讨论
4.1 数据集和评估
突触多器官分割数据集1 我们在MICCAI2015多地图集腹部标签挑战中使用了30张腹部CT扫描，共使用了3779张轴向对比度增强的腹部临床CT图像。每个CT体积由85∼198片512×512像素组成，体素空间分辨率为（[0.54∼0.54]×[0.98∼0.98]×[2.5和5.0]）mm3。 在[5]之后，我们报告了平均DSC和平均Hausdorff距离(HD)8个腹部器官（主动脉、胆囊、脾、左肾、右肾、肝、胰腺、脾、胃，随机分裂18例（2212片），12例验证。
自动的心脏诊断挑战2 ACDC挑战收集从MRI扫描仪获得的不同患者的检查。电影先生图像是通过呼吸屏住而获得的，一系列的短轴切片覆盖了心脏从心底到左心室的顶端，切片厚度为5到8毫米。短轴平面内空间分辨率从0.83上升到1.75mm2/像素。每次患者扫描均会人工标注左心室(LV)、右心室(RV)和心肌(MYO)的地面真相。我们报告了平均DSC，随机分割了70个训练案例（1930个轴向切片），10例用于验证，20例用于测试。
4.2所有实验的实施细节
我们应用简单的数据增强，例如随机旋转和翻转。 对于纯变压器编码器，我们简单地采用VIT[4]与12个变压器层。 对于混合编码器的设计，我们结合ResNet-50[6]和VIT，表示为“R50-VIT”，通过本文。 所有变压器骨干(即VIT)和ResNet-50(表示为“R-50”)都在图像网[2]上进行了预先培训。 输入分辨率和补丁大小P设置为224×224和16，除非另有规定。 因此，我们需要在CUP中连续级联四个2×上采样块，以达到完全分辨率。 对于模型，使用学习速率0.01、动量0.9和重量衰减1e-4的SGD优化器进行训练。 默认批处理大小为24，默认训练迭代次数分别为ACDC数据集20k和Synapse数据集14k。 所有实验都是使用单个NvidiaRTX2080Ti GPU进行的。在[17,19]之后，分层推断所有的三维卷，并将预测的二维切片叠加在一起，以重建三维预测以进行评估。
4.3与最先进的比较
我们通过对突触多器官分割数据集进行了主要实验，与我们的四个技术进行比较：1)V网[9]；2)DARR[5]；3)u网[12]和4)[13]。为了证明杯解码器的有效性，我们使用ViT[4]作为编码器，并分别使用朴素上采样（“无”）和CUP作为解码器比较结果；为了验证我们的混合编码器设计的有效性，我们使用CUP作为解码器，并分别使用ViT和R50-ViT作为编码器进行比较。为了公平地与动力基线(R50-ViT-CUP)和我们的跨网络进行比较，我们还用图像网络预训练的[12]和[10]的原始编码器。DSC和平均行李间隔距离(mm)的结果见表1。
首先，我们可以看到，与VIT-None相比，VIT-CUP在平均DSC和Hausdorff距离方面分别有6.36%和3.50mm的改善。 这一改进表明，我们的CUP设计提出了比直接上采样更好的解码策略。 同样，与VIT-CUP相比，R50-VIT-CUP在DSC和Hausdorff距离上也增加了3.43%和3.24mm，这表明了我们的混合编码器的有效性。 基于R50-Vi T-CUP，我们的Trans UNet也配备了跳接，在基于变压器的模型的不同变体中取得了最好的效果。 其次，表1还显示，拟议的TransUnet比以前的艺术有了显著的改进，例如，考虑到平均DSC，性能增益从1.91%到8.67%不等。 特别是，直接应用变压器进行多器官分割得到了合理的结果(67.86%DSC用于VIT-CUP)，但不能与U-Net或ATTUNET的性能相匹配。这是因为变形金刚可以很好地捕捉到有利于分类任务的高级语义，但缺乏低级的线索来分割精细形状的奖牌
4.4分析研究
为了彻底评估所提出的跨式单元仪框架，并验证在不同设置下的性能，我们进行了各种消融研究，其中：
1）跳过连接的次数；
2）输入分辨率；
3）序列长度和补丁大小，
4）模型缩放。
跳过连接的数量。如上所述，集成u网络跳过连接通过恢复底层空间信息有助于增强更精细的分割细节。此消融的目的是测试在跨单元中添加不同数量的跳过连接的影响。通过改变跳过连接数为0（R50-ViT-CUP）/1/3，所有8个检测机构的平均DSC的分割性能总结见图2。请注意，在“1跳过”设置中，我们仅在1/4分辨率比例上添加跳过连接。我们可以看到，添加更多的跳过连接通常会导致更好的分割性能。最佳平均DSC和HD通过插入除输出层以外的所有1/2、1/4和1/8分辨率步骤的跳过连接来实现，如图1所示）。因此，我们对跨联合国采用了这种配置。同样值得一提的是，较小器官（即主动脉、胆囊、肾、胰腺）的性能增加更为明显
作为一项有趣的研究，我们将附加变压器应用于跳过连接中，类似于[13]，并发现这种新型的跳过连接甚至可以进一步提高分割性能。由于GPU内存约束，我们在1/8分辨率比例的跳过连接中使用了一个轻型变压器，同时保持其他两个跳过连接不变。因此，这种简单的改变导致性能提高了1.4%的DSC。
对输入解析方法的影响。跨UN集的默认输入分辨率为224×224。在这里，我们还提供了在高分辨率512×512上的训练结果，如表2所示。当使用512×512作为输入时，我们保持相同的补丁大小（即16），这就导致了一个近似的结果,5×变压器的序列长度。正如[4]所指出的，增加有效序列长度会显示出稳健的改进。对于跨单元，将分辨率尺度从224×224改为512×512，平均DSC提高6.88%，代价是更大的计算代价。因此，考虑到计算成本，本文中所有的实验比较都以默认分辨率为224×224进行，以证明跨单元的有效性。
对补丁程序大小/序列长度的影响。我们还研究了补丁大小对跨单元网的影响。研究结果的汇总情况见表3。结果表明，用较小的斑块大小，通常可以获得较高的分割性能。请注意，变压器的序列长度与补丁大小的平方成反比（例如，补丁大小16对应于序列长度196，而补丁大小32具有较短的序列长度49），因此减小补丁大小（或增加有效序列长度）显示出稳健的改进，因为变压器为较长的输入序列编码每个元素之间更复杂的依赖关系。在[4]中设置之后，我们在本文中使用16×16作为默认补丁大小。
模型缩放。 最后但并非最不重要的是，我们提供了不同型号的TransUNET消融研究。 特别是，我们研究了两种不同的跨联合国网络配置，“基地”和“大”模型。 对于“基本”模型，隐藏大小D、层数、MLP大小和头数分别设置为12、768、3072和12，而对于“大”模型的超参数分别为24、1024、4096和16。 从表4中我们得出结论，较大的模型会导致更好的性能。 考虑到计算成本，我们对所有实验采用了“基本”模型。
4.5可视化
我们在Synapse数据集上提供定性比较结果，如图3所示。 可以看出：1)基于CNN的纯方法U-Net和AttnUnet更有可能对器官进行过度分割或欠分割(例如，在第二行中，脾脏被AttnUnet过度分割，而被UNet不足分割)，这表明基于变压器的模型，例如我们的TransUnet或R50-ViT-CUP具有更强的编码全局上下文和区分语义的能力。 2)第一行的结果表明，与其他方法相比，我们的反式UNet预测的假阳性较少，这表明反式UNet在抑制这些噪声预测方面比其他方法更有利。 3)为了在基于Transformer的模型中进行比较，我们可以观察到R50-ViT-CUP的预测往往比TransUnet关于边界和形状的预测更粗糙(例如，第二个胰腺的预测第1行）。此外，在第三行，跨网正确地预测了左右肾，而R50-ViT-杯错误地填补了左肾的内孔。这些观察结果表明，跨单元能够更好的分割和保留详细的形状信息。原因是跨单元网享有高层次全局上下文信息和低层次细节的好处，而R50-Vi-T-CUP完全依赖于高级语义特征。这再次验证了我们最初的直觉，即将u网样的跳过连接集成到变压器设计中，以实现精确的定位。
4.6泛化到其他数据集
为了显示我们的跨数据集的泛化能力，我们进一步评估了其他成像模式，即旨在自动心脏分割的Mr数据集ACDC。我们观察到跨UNT比纯基于CNN的方法(R50-UNT和R50连接UNT)和其他基于转换的基线(ViT-CUP和R50-ViT-CUP)有一致的改进，这与之前在突触CT数据集上的结果相似。
5结论
Transformers是一种具有较强固有自我注意机制的结构。 本文首次研究了变形金刚在普通医学图像分割中的应用。 为了充分利用变形金刚的力量，提出了TransUNet，它不仅通过将图像特征视为序列来编码强大的全局上下文，而且通过u形混合结构设计很好地利用了低级别的CNN特征。 作为一种替代基于FCN的医学图像分割方法的替代框架，TransUnet比各种竞争方法(包括基于CNN的自我关注方法)具有更好的性能。 致谢。 这项工作得到了Lustgarten胰腺癌研究基金会的支持。