[深度学习论文笔记]UNETR: Transformers for 3D Medical Image Segmentation

UNETR: Transformers for 3D Medical Image Segmentation
UNETR：用于三维医学图像分割的Transformer

Published: Oct 2021
Published in: IEEE Winter Conference on Applications of Computer Vision (WACV) 2022
论文：https://arxiv.org/abs/2103.10504
代码：https://monai.io/research/unetr

摘要：
  过去十年以来，具有收缩路径和扩展路径（编码器和解码器）的全卷积神经网络（FCNNs）在各种医学图像分割应用中表现突出。在这些体系结构中，编码器通过学习全局和局部特征以及上下文表示成为一个不可或缺的角色，这些特征和上下文表示可以被解码器用于语义输出预测。尽管它们取得了成功，但卷积层在FCNN中的局限性，限制了在此类网络中学习远程空间相关性的能力。受自然语言处理（NLP）在远程序列学习中的最新成功的启发（Transformers），作者将体积（3D）医学图像分割的任务重新设计为序列到序列的预测问题。介绍了一种称为UNEt TRansformers（UNETR）的新颖架构，该架构利用纯Transformers作为编码器来学习输入量的序列表示并有效地捕获全局多尺度信息。同时也遵循了编码器和解码器的成功的“U型”网络设计,Transformers编码器通过不同分辨率的跳跃连接直接连接到解码器，以计算最终的语义分段输出。使用多器官分割的多图谱标记(BTCV)数据集、医学分割十项全能（MSD）数据集广泛验证了提出的模型在不同成像方式（即MR和CT）上对体积脑肿瘤和脾脏分割任务的性能，并且结果始终证明了良好的性能。

问题动机：
  图像分割在定量医学图像分析中起着不可或缺的作用，因为它往往是解剖结构分析的第一步。自从深度学习出现以来，FCNNS，特别是“U型”编解码架构已经在各种医学语义分割任务中取得了最先进的结果。在一个典型的U-Net架构中，编码器负责通过逐步降低提取特征的采样来学习全局上下文表示，而解码器负责将提取的表示采样到输入分辨率，以进行像素/体素的语义预测。此外，跳过连接合并编码器的输出和解码器在不同的分辨率，因此允许恢复在降采样期间丢失的空间信息。
  虽然这种基于FCN的方法具有强大的表示学习能力，但它们在远程依赖学习中的表现，仅限于它们的局部接受域。因此，这种在多尺度信息捕获方面的缺陷导致了对不同形状和尺度的结构(如不同大小的脑损伤)的次优分割存在不足。一些研究已经使用atrous convolutional layers来扩大感受域。然而，卷积层中接受域的局部性仍然限制了它们的学习能力到相对较小的区域。将自注意模块与卷积层相结合，提高了非局部建模能力。
  在自然语言处理(NLP)中，基于Transformers的模型在各种任务中实现了最先进的基准。Transformers的自注意机制允许动态突出词序列的重要特征。此外，在计算机视觉中，使用Transformers作为主干编码器是有益的，因为它们具有建模远程依赖和捕获全局上下文的强大能力。具体来说，与卷积的局部公式不同，Transformers将图像编码为1维 patch嵌入序列，并利用自注意模块学习从隐藏层计算出的值的加权和。因此，这种灵活的提法可以有效地学习长距离信息。此外，视觉Transformers (ViT)及其变体在学习文本前任务方面表现出了出色的能力，这些任务可以转移到下游应用任务。
  在这项工作中，作者提出利用Transformers的能力进行体积医学图像分割，并引入了一种新的结构称为UNEt TRansformers (UNETR)。特别的，作者将三维分割任务重新表述为一维序列对序列的预测问题，并使用Transformers作为编码器从嵌入的输入补丁中学习上下文信息。从Transformers编码器中提取的表示通过多个分辨率的跳过连接与基于CNN的解码器合并，以预测分割输出。提出的框架使用基于CNN的解码器，而不是在解码器中也使用Transformers。这是由于尽管Transformers具有很强的全局信息学习能力，但却不能正确地捕获局部信息。
  作者使用Beyond the Cranial Vault (BTCV)和Medical segmentation Decathlon (MSD)数据集验证了方法在3D CT和MRI分割任务中的有效性。在BTCV数据集中，UNETR在其排行榜上的标准和自由竞赛部分都实现了新的最先进的表现。在MSD数据集中，UNETR在脑肿瘤和脾脏分割任务方面的表现优于最先进的方法。

主要贡献如下:
  1、提出了一种新的基于Transformers的医学图像分割模型。
  2、提出了一种新颖的架构，其中:(1) Transformers编码器直接利用嵌入的3D体块有效地捕捉远程依赖;(2)跳跃式连接解码器结合提取的不同分辨率的表示，并预测分割输出。
  3、在两个公共数据集(BTCV和MSD)上验证了提出的模型对不同体积分割任务的有效性。UNETR在BTCV数据集排行榜上实现了新的最先进的性能，并在MSD数据集上优于竞争方法。

思路来源：
  CNN-based分割网络：自从引入开创性的U-Net以来，基于CNN的网络已经在各种2D和3D各种医学图像分割任务上取得了最先进的结果。对于以体积为单位的分割，有时使用三平面结构来组合每个体素的三视图切片，也称为2.5D方法。相比之下，3D方法直接利用由一系列二维切片或形式表示的完整的体积图像。采用不同尺寸的直观方法之后，采用多扫描、多路径模型来捕获图像的下采样特征。此外，为了充分利用三维环境和解决计算资源有限的问题，研究人员还研究了层次结构框架。
  有人提出在多个尺度或组合框架上提取特征。提出的多尺度框架来获得胰腺分割中不同分辨率的信息减少了空间上下文和低分辨率条件下的问题，为三维医学图像的多层次分割提供了先导性研究。尽管这些网络取得了成功，但其局限性在于它们在学习全局上下文和远程空间依赖方面的性能较差，这可能会严重影响具有挑战性任务的分割性能。

  Vision Transformers：视觉Transformers最近获得了计算机视觉任务的关注。Dosovitskiy等人通过对纯Transformers进行大规模的预训练和微调，展示了图像分类数据的最先进性能。在目标检测中，端到端基于Transformers的模型在几个基准测试中表现突出。最近提出了具有不同分辨率和空间嵌入的层次视觉Transformers。这些方法逐渐降低了Transformers层特征的分辨率，并利用了次采样注意模块。与这些方法不同的是，UNETR编码器中表示的大小在所有Transformers层中保持固定，使用反卷积和卷积操作来改变提取特征的分辨率。
  最近，人们提出了多种方法，探索使用基于Transformers的模型进行二维图像分割的可能性。Zheng等人引入了SETR模型，其中预训练的Transformers编码器与不同的基于CNN的解码器被提出用于语义分割任务。Chenet al.提出了一种方法，通过在U-Net体系结构的瓶颈中使用Transformers作为额外的一层来实现多器官分割。Zhang等人建议将CNN网络和Transformers分开使用，并融合它们的输出。Valanarasu等人提出了一种基于Transformers的轴向注意机制用于二维医学图像分割。作者的模型与这些工作有关键的区别:
  (1) UNETR是为3D分割而量身定制的，并直接利用体积数据;
  (2) UNETR使用Transformers作为分割网络的主要编码器，并通过跳过连接将其直接连接到解码器，而不是使用它作为分割网络中的注意层；
  (3) UNETR不依赖骨干CNN来生成输入序列，而是直接利用标记化的patch。

  对于3D医学图像分割，Xie等人提出了一种框架，利用骨干CNN进行特征提取，利用Transformers处理编码表示，利用CNN解码器预测分割输出。类似地，Wanget al.提出在3D编解码器CNN的瓶颈处使用Transformers来完成脑肿瘤语义分割的任务。与这些方法不同的是，作者的方法通过使用跳跃连接直接将Transformers的编码表示连接到解码器。

模型方法:
  在图2中展示了所提议的模型的概述。UNETR采用收缩-扩展模式，由一堆Transformer组成，编码器通过跳跃连接到解码器。Transformer工作在输入嵌入的一维序列上，这是NLP中常用的方法。同样，将三维输入体积 ，分辨率(H,W,D)和C输入通道划分为平坦的均匀非重叠斑块 ，创建一个一维序列，其中(P, P, P)表示每个patch的分辨率， 为序列的长度。



  图2：UNETR架构概述。将一个三维输入体(如MRI图像的C= 4通道)划分为一系列均匀的不重叠的patch，并使用线性层投影到嵌入空间中。该序列与位置嵌入一起添加，并用作Transformer模型的输入。Transformer中不同层的编码表示通过跳跃连接提取并与解码器合并，以预测最终分割。

  图1所示：UNETR的概述。提出的模型包括一个Transformers编码器，它直接利用3D patch，并通过跳跃连接 连接到一个基于CNN的解码器。

损失函数：
  损失函数是soft-dice损失和交叉熵损失的组合，它可以按体素方式计算：

其中I为体素的数量;J为类数;Yi,j Gi,j分别表示j类在体素i上的概率输出和one-hot编码ground truth。

实验与结果：
数据集：
  为了验证方法的有效性，利用BTCV和MSD数据集在CT和MRI成像模式中进行三种不同的分割任务。

  BTCV (CT): BTCV数据集包括30名进行腹部CT扫描的受试者，其中13个器官由标注人员在范德比尔特大学医学中心临床放射科医生的监督下，进行标注。每次CT扫描均在门静脉期进行造影增强，共80,225片 512×512 pixels，层厚1 ~ 6mm。通过将强度在[-1000,1000]HU到[0,1]的范围内归一化，分别对每个体积进行预处理。在预处理过程中，将所有图像重新采样到各向同性体素间距1.0mm。多器官分割问题是一个13类1通道输入的分割任务。

  MSD(MRI/CT):对于脑肿瘤的分割任务，使用整个训练集484个多模态多位点MRI数据(FLAIR, T1w, T1gd, T2w)，并带有胶质瘤分割坏死/活性肿瘤和水肿的ground truth标签进行模型训练。本任务中MRI图像体素间距为1.0 × 1.0 × 1.0 mm3。对体素强度进行z-score归一化预处理。脑肿瘤分割问题是一个具有4通道输入的3类分割任务。

  在脾脏分割任务中，使用了41个带有脾脏体注释的CT体积。任务9中卷的分辨率/间距范围为0.613×0.613×1.50 mm3 ~ 0.977×0.977×8.0 mm3。在预处理过程中，将所有体积重新采样到各向同性体素间距为1.0 mm。将图像的体素强度按整体前景强度的第5和第95个百分位归一化到[0,1]范围。脾脏分割是一个具有1通道输入的二进制分割任务。对于多器官和脾脏分割任务，我们对输入的图像随机采样，体积大小为[96,96,96]。对于脑分割任务，对输入的图像随机采样，体积大小为[128,128,128]。在所有的实验中，前景/背景的随机斑块以1:1的比例采样。

评价指标：
  使用Dice评分和95% Hausdorff Distance (HD)来评估分割的准确性。对于给定的语义类，让Gi和 Pi表示体素i的地面真值和预测值，G’和P’分别表示ground真值和预测面点集。
Dice分数和HD指标被定义为:

95% HD使用地面真实值和预测表面点集之间距离的第95个百分位数。因此，在计算HD时，一个非常小的异常子集的影响是最小的。

实现细节：
  在pytorch和MONAI中实现UNETR。该模型使用NVIDIA DGX-1服务器进行训练。所有模型都在批大小为6的情况下进行训练，使用初始学习率为0.0001的AdamW优化器进行20000次迭代。对于指定的批处理大小，20000次迭代的平均训练时间为10小时。基于Transformers的编码器遵循ViT-B16架构，L= 12层，嵌入大小k = 768。使用的补丁分辨率为16×16×16。为了进行推断，使用滑动窗口方法，在相邻补丁之间的重叠部分为0.5，分辨率与第4.1节中指定的相同。没有对Transformers主干使用任何预先训练的权重(例如ImageNet上的ViT)，因为它没有显示出任何性能改进。对于BTCV数据集，已经在其排行榜的标准和免费竞赛中评估了模型和其他基线。自由竞赛使用了来自同一队列的额外数据，使训练案例增加到80卷。对于所有的实验，采用了5倍交叉验证，比例为95:5。此外，使用了数据增强策略，如随机旋转90度、180度和270度，随机翻转轴位、矢状位和冠状位以及随机尺度和偏移强度。使用集成来融合四个不同的五重交叉验证模型的输出。对于MSD数据集中的大脑和脾脏分割任务，将数据按照80:15:5的比例分成训练、验证和测试三部分。

定量评估：
  在BTCV排行榜上，UNETR在标准比赛和免费比赛中表现优于最先进的方法。如表1所示，在自由竞赛中，UNETR的整体平均Dice得分为0.899，比排名第二、第三和第四的方法分别高出1.238%、1.696%和5.269%。

  表1：BTCV测试集的分割性能的定量比较。顶部和底部分别代表标准比赛和自由比赛的基准。作者的方法与当前最先进的模型进行了比较。

  在标准竞赛中，比较了UNETR与CNN和基于Transformers的基线的性能。UNETR实现了一个新的最先进的性能与平均骰子得分85.3%在所有器官。在脾、肝、胃等大器官上，方法在Dice得分上分别比次优基线高出1.043%、0.830%和2.125%。此外，在小器官的分割中，在胆囊和肾上腺的Dice评分方面，方法显著优于第二最佳基线的6.382%和6.772%。
  在表2中，在MSD数据集上比较了UNETR与CNN以及基于Transformers的方法在脑肿瘤和脾脏分割任务中的性能。在脑分割方面，UNETR在所有语义类上的平均表现比最接近的基线高出1.5%。特别是，UNETR在分割肿瘤核心(TC)分区域方面表现得相当好。同样，在脾脏分割方面，UNETR在Dice评分方面的表现比最佳竞争方法高出至少1.0%。

  表2：定量比较了MSD数据集在脑肿瘤和脾脏分割任务中的分割性能。WT、ET和TC分别表示全肿瘤、增强肿瘤和肿瘤核心子区域。

定性结果：
  多器官分割的定性比较如图3所示。UNETR提高了腹部器官的分割性能。模型学习远程依赖的能力在第3行(从上面开始)很明显，其中nnUNet混淆了肝脏和胃组织，而UNETR成功地描绘了这些器官的边界。在图3中，第2行和第4行显示肾脏和肾上腺对周围组织的清晰检测，这表明UNETR捕获了更好的空间上下文。与基于二维Transformers的模型相比，UNETR展示了更高的边界分割精度，因为它准确地识别了肾脏和脾脏之间的边界。可见胆囊在第二排，肝脏和胃在第三排，门静脉对肝脏在第五排。在图4中，给出了在MSD数据集上对脑肿瘤分割的定性分割比较。具体来说，模型在捕捉肿瘤的细粒度细节方面表现得更好。

  图3：BTCV交叉验证中不同基线的定性比较。第一行为完整代表性CT切片。展示了4个放大的受试者(第2至5行)，方法显示了肾脏和脾脏(第2行)、胰腺和肾上腺(第3行)、胆囊(第4行)和门静脉(第5行)的分割效果。每个样本上显示了受试者的平均Dice得分。

  图4：UNETR有效地捕捉细分输出的细粒度细节。
  在所有数据集的实验表明，UNETR的性能优于CNN和基于Transformers的分割模型。具体来说，UNETR通过捕获全局和局部依赖实现了更好的分割精度。在定性比较中，这一点在各种案例中得到了说明，在这些案例中，UNETR有效地捕捉了长期依赖关系(例如图3中对胰腺尾部的精确分割)。
  此外，UNETR在BTCV排行榜上的细分表现显示了新的最先进的基准，并验证了其有效性。特别是对于小的解剖，UNETR优于CNN和基于Transformers的模型。虽然3D模型已经证明了对小器官(如胆囊、肾上腺)的高分割精度，但UNETR仍能以显著优势领先于最佳竞争模型(见表1)。
这在图3中也可以观察到，其中UNETR对左肾上腺和右肾上腺的分割精度明显更好，而且UNETR是唯一正确检测肾上腺分支的模型。对于较具挑战性的组织，如第4行胆囊和第5行门静脉，它们与周围肝组织的对比度较低，UNETR仍然能够分割出清晰的连接边界。

消融实验：
  解码器的选择：在表3中，通过比较UNETR和其他解码器架构在两个MRI和CT模式的代表性分割任务上的性能来评估解码器的有效性。在这些实验中，使用UNETR的编码器，但用朴素向上采样(NUP)、渐进向上采样(PUP)和多尺度聚合(MLA)的3D对应版本替换解码器。观察到这些解码器架构产生次优性能，尽管MLA略微优于NUP和PUP。对于脑肿瘤的分割，UNETR比MLA、PUP和NUP解码器的平均准确率分别高出2.7%、4.3%和7.5%。同样，对于脾脏分割，UNETR外形成的MLA、PUP和NUP分别为1.4%、2.3%和3.2%。

  表3：解码器架构对分割性能的影响。（Naive UpSampling、Progressive UpSampling and Multi-scale Aggregation）

  Patch Resolution：较低的输入补丁分辨率导致较大的序列长度，因此内存消耗更高，因为它与分辨率的立方成反比。如表4所示，实验表明，降低分辨率可以持续提高性能。具体来说，将patch分辨率从32降低到16，在脾脏和大脑分割任务的平均Dice得分方面分别提高了1.1%和0.8%的性能。由于内存的限制，没有进行低分辨率的实验。

  表4：patch分辨率对分割性能的影响。

  模型与计算复杂度：在表5中，给出了BTCV基准中模型的FLOPs次数、参数和平均推理时间。浮点数和推断时间是基于96×96×96的输入大小并使用滑动窗口方法计算的。根据基准，UNETR是一个中等大小的模型，具有92.58M参数和41.19G浮点数。相比而言，其他基于Transformers的方法，如CoTr、TransUNet和SETR，分别有46.51M、96.07M和86.03M参数和399.21G、48.24G和43.49G FLOPs。UNETR显示了类似的模型复杂性，同时在BTCV基准中大大优于这些模型。nnUNet和ASPP基于cnn的分割模型分别有19.07M和47.92M参数，412.65G和44.87G FLOPs。同样，UNETR优于这些基于cnn的模型，同时具有中等的模型复杂性。此外，UNETR的平均推理时间仅次于nnUNet，而且明显快于基于Transformers的模型，如SETR， TransUNet和CoTr。

  表5所示：BTCV实验中各模型的参数数、次数和平均推断时间的比较。

结论：
  本文介绍了一种新颖的基于Transformer的结构，称为UNETR，用于体积医学图像的语义分割，将该任务重新定义为一个一维序列到序列的预测问题。作者建议使用Transformers编码器来提高模型的能力，以学习远程依赖关系，并在多个尺度上有效地捕获全局上下文表示。
  验证了UNETR在CT和MRI不同体积分割任务中的有效性。在BTCV多器官分割排行榜上，UNETR竞赛中取得了最新的水平表现，并在MSD数据集上优于脑肿瘤和脾脏分割的竞争方法。该方法可作为医学图像分析中一类新的基于Transformers的分割模型的基础。