UNETR: Transformers for 3D Medical Image Segmentation

Posted 2021-12-23 Donogh Liu

UNETR: Transformers for 3D Medical Image Segmentation

[代码地址]：(https://monai.io/research/unetr)

摘要

自过去十年以来，具有收缩和扩展路径的完全卷积神经网络（FCNN）在大多数医学图像分割应用中显示出突出的地位。在FCNNs中，编码器通过学习全局和局部特征以及可由解码器用于语义输出预测的上下文表示来发挥不可或缺的作用。尽管它们取得了成功，但FCNN中卷积层的局部性限制了学习长程空间依赖性的能力。受自然语言处理转换器（NLP）最近在远程序列学习中取得成功的启发，我们将三维医学图像分割任务重新定义为序列到序列的预测问题。我们介绍了一种新的架构，称为UNEt Transformers（UNETR），它利用一个转换器作为编码器来学习输入体积的序列表示，并有效地捕获全局多尺度信息，同时也遵循编码器和解码器成功的“U形”网络设计。transformer编码器通过不同分辨率的跳过连接直接连接到解码器，以计算最终的语义分段输出。这是第一个为3维图像提出的Transformer分割模型。并在多器官分割的BTCV挑战赛，以及脑肿瘤和脾脏分割任务的医学分割十项全能（MSD）数据集上展现了最先进的性能。

方法比较与CNN

1.在典型的U形网中[36]在体系结构中，编码器负责通过逐渐对提取的特征进行下采样来学习全局上下文表示，而解码器则将提取的表示向上采样到输入分辨率，以进行像素/体素语义预测。此外，跳过连接将编码器的输出与解码器在不同位置合并租用分辨率，因此允许恢复在下采样期间丢失的空间信息。但是但它们在学习长距离依赖性方面的性能仅限于其局部感受野。因此，这种在获取多尺度信息方面的缺陷导致了对具有不同形状和尺度的结构（例如不同大小的脑损伤）的次优分割。(尽管提出了一些解决方法：空洞卷积)，但任然限制了相对小区域的的学习能力。对于卷积网络的一个共同的特点，这些网络的一个局限性是它们在学习全局上下文和长距离空间依赖性方面的性能较差，这会严重影响具有挑战性任务的分割性能。
2.计算机视觉中，使用变压器(Transformer)作为主干编码器是有益的，因为它们具有很强的建模长期依赖性和捕获全局上下文的能力。具体而言，与卷积的局部公式不同，transformers将图像编码为一维面片嵌入序列，并利用自我注意模块学习从隐藏层计算的值的加权和。因此，这种灵活的公式可以有效地学习远程信息。

特点：（1） UNETR专门用于三维分割，并直接利用体积数据；（2） UNETR将变压器用作分段网络的主编码器，并通过跳过连接将其直接连接到解码器，而不是将其用作分段网络内的注意层（3）。UNETR不依赖主干CNN生成输入序列，并直接利用标记化补丁。

方法

网络结构

如图：一个3D输入体（例如，MRI图像的C=4通道）被分成一系列均匀的非重叠贴片，并使用线性层投影到嵌入空间。该序列与位置嵌入一起添加，并用作变压器模型的输入。提取变压器中不同层的编码表示，并通过跳过连接与解码器合并，以预测最终分段。Output sizes 的patch = 16，embedding size k=768。

与NLP中常用的一样，变压器在1D输入嵌入序列上运行。类似地，我们创建一个3D体素的一维输入序列，3D体素的输入为（H,W,D,C),将其平坦化非重叠的(N*(P^3.C))。N = (HWD)/P^3
,为序列长度。其在运算的过程和transformer的计算类似。其实就是使用transformer作为一个编码结构提取特征，学习特征的过程。

损失函数

作者提出使用交叉熵和骰子系数的组合损失函数。

实验结果

1. BTCV（CT）：BTCV数据集[26]由30名腹部CT扫描的受试者组成，其中13个器官在V anderbilt大学医学中心临床放射科医生的监督下由译员注释。每次CT扫描均在门静脉期进行对比增强，由80至225层组成，512×512像素，层厚为1至6mm。通过对[-10001000]HU至[0,1]范围内的强度进行归一化，对每个体积分别进行了预处理。在预处理时所有图像都将重采样为1mm的各向同性体素间距。将多器官分割问题描述为一个单通道输入的13类分割任务。
   

2. MSD（MRI/CT）：对于脑肿瘤分割任务，使用484个多模式多部位MRI数据（FLAIR、T1w、T1gd、T2w）的整个训练集，以及胶质瘤分割坏死/活动性肿瘤和水肿的基本真实标签进行模型训练。本任务中MRI图像的体素间距为1。0×1. 0×1. 0立方米。体素强度通过z分数归一化进行预处理。将脑肿瘤分割问题描述为一个具有4通道输入的3类分割任务。对于脾脏分割任务，使用了41个带有脾体注释的CT体积。任务9中卷的分辨率/间距范围为0。613×0. 613×1. 50mm至0。977×0. 977×8. 0mm。将所有体积重新采样到各向同性体素间距1.0mm在预处理期间。根据总前景强度的第5百分位和第95百分位，将图像的体素强度标准化为范围[0,1]。脾脏分割是一个单通道输入的二进制分割任务。对于多器官和脾脏分割任务，我们随机抽取体积大小为[96,96,96]的输入图像。对于大脑分割任务，我们随机抽取体积大小为[128128]的输入图像。对于所有实验，前景/背景的随机斑块都是以1:1的比例采样的。
   
   

讨论

本文介绍了一种新的基于变换器的结构，称为UNETR，通过将此任务重新表述为一维序列到序列预测问题，对体积医学图像进行语义分割。我们建议使用transformer编码器来提高模型学习长期依赖关系的能力，并在多个尺度上有效地捕获全局上下文表示。我们验证了UNETR在CT和MRI模式中对不同体积分割任务的有效性。UNETR在BTCV多器官分割排行榜上的标准和免费竞赛中都取得了最新的水平，并且在MSD数据集上的脑肿瘤和脾脏分割方面的表现优于竞争对手。总之，UNETR已显示出有效学习医学图像中关键解剖关系的潜力。所提出的方法可以为医学图像分析中一类新的基于变压器的分割模型奠定基础。