Extending nn-UNet for brain tumor Segmentation 扩展nn-UNet用于脑肿瘤分割

Extending nn-UNet for brain tumor Segmentation
扩展nn-UNet用于脑肿瘤分割

Pattern Recognition on November 01, 2021
论文：https://arxiv.org/abs/2112.04653
code：https://github.com/rixez/Brats21_KAIST_MRI_Lab
https://hub.docker.com/r/rixez/brats21nnunet

摘要：
  脑肿瘤的分割对胶质瘤的诊断和预后至关重要。脑肿瘤分割的挑战继续为开发自动算法来执行任务提供了大量的数据来源。开发的方法基于nn-UNet，这是去年比赛的获奖作品。试验了几种改进，包括使用更大的网络，①用组归一化替换批量归一化，以及②在解码器中使用轴向注意力。内部的5倍交叉验证以及组织者的在线评估都显示了方法的有效性，与基线相比，量化指标略有改进。在未见测试数据的最终排名中，提出的模型赢得了第一名。获奖作品的代码、预训练的权重和docker图像都是公开的。

问题动机：
  从磁共振(MR)图像中分割脑肿瘤是脑肿瘤治疗的一个重要步骤，使临床医生能够确定肿瘤的位置、范围和类型。这不仅有助于初步诊断，也有助于管理和监测治疗进展。考虑到这项任务的重要性，精确描绘肿瘤及其子区域通常由经验丰富的神经放射学家手工完成。这是一个冗长而耗时的过程，需要大量的时间和专业知识，特别是当病人体积大，图像是多参数的不同对比度，肿瘤是异质性的时候。标签过程也受制于内部和内部评级者的差异性，需要对标签和分割的解释达成共识，并增加了额外的复杂性层。自动或计算机辅助分割算法有潜力解决这些缺点，因为它可以降低标签过程的劳动强度，并在不同情况下保持一致。然而，为了开发这些算法，需要有足够的高质量的注释数据来达到临床目的。
  对脑肿瘤的自动分割的最初尝试依赖于手工制作的特征工程和传统的机器学习方法，如基于atlas，决策森林，条件随机场。随着现代图形处理单元(GPU)计算能力的提高、算法效率的提高以及训练数据可用性的提高，深度学习的日益普及，在计算机视觉等几个领域，传统的机器学习方法已经慢慢被深度神经网络所取代，自然语言处理，或计算生物学。在BraTS竞赛的背景下，自2014年以来，深度学习算法被用于肿瘤分割，它们是近年来大多数参赛作品的选择算法。事实上，最近4个比赛的获胜者都使用了深度神经网络，证明了这种方法在获得更多数据时的优越性能。Kamnitsas等人结合来自不同3D卷积网络(DeepMedic、FCN和U-Net)的预测，提出了多模型和架构集成(EMMA)。Myronenko结合了一个3D U-Net和一个额外的变分解码器分支，以提供额外的监督和正则化的编码器分支。Jiang等人训练了一个两阶段级联U-Net，第一阶段训练产生粗分割掩模，第二阶段训练提炼第一阶段的输出。Isensee等人利用nnU-Net，这是一个自动适应U-Net到特定数据集的自配置框架，并通过使用一些brats特定的优化，在对传统3D U-Net进行最小修改的情况下显示了鲁棒性能。
  为了参加比赛，扩展了去年获胜者提出的nnU-Net框架，增加了几个组件。由于nnU-Net易于适应新的数据集以及完全开放源代码的代码和模型，nnU-Net可以作为进一步实验的优秀基准。为今年的比赛所做的修改是使用组归一化而不是批量归一化，为U-Net使用不对称的大编码器，并在解码器中使用轴向注意。交叉验证训练数据和排行榜中未见验证数据的实验结果表明，这些修改对网络性能的影响。

思路来源：
  来自2000名患者的多参数MRI扫描被用于BraTS2021，其中1251人被提供分割标签给参与者来开发他们的算法，其中219人在验证阶段被用于公共排行榜，其余530个案例用于私人排行榜和参与者的最终排名。MRI扫描有4种对比:原生t1加权图像、对比后t1加权(T1Gd)、T2加权和T2流体衰减反转恢复(T2- flair)。注释由一到四位专业人员手工完成，最终得到有经验的神经放射学家的验证。这些标签包括gd增强肿瘤(ET)、瘤周水肿/侵袭组织(ED)和坏死肿瘤核心(NCR)区域。所有MRI扫描均通过同解剖模板配准、各向同性1mm3分辨率插值和颅骨剥离进行预处理。所有MRI扫描和相关标签的图像大小为240×240×155。图1显示了这四种对比与分割的代表性切片。在输入网络之前，对提供的数据进行了进一步的处理。为了减少计算量，将体积裁剪为非零体素。由于MR图像的强度是定性的，因此根据其均值和标准差对体素进行归一化。

baseline nnU-Net：
  Isensee等人的nnU-Net是2020年BraTS的获奖作品。其核心是一个在128× 128×128大小的patch上运行的3D U-Net。该网络具有编码-解码器结构，并带有跳跃式连接，将两条路径连接起来。该编码器由5个相同分辨率的卷积层组成，具有跨步卷积下采样功能。该解码器遵循相同的结构，使用转置卷积上采样和卷积操作在同一级别上的编码器分支的串联跳跃特征。

  图1：一个受试者的代表性轴向切片(BraTS2021 00001)显示了T1图像上3个肿瘤子区域的四种不同对比度和重叠的分割标签。绿色:瘤周水肿/侵犯组织，红色:坏死肿瘤核心，黄色:gd增强肿瘤。
  每次卷积操作后使用批量归一化和Leaky ReLU (lReLU)。mpMRI体积被连接并作为4通道输入。nnU-Net提出训练要求:而不是预测三个互斥肿瘤分区,如提供细分标签,网络预测而不是三个重叠的地区加强肿瘤(ET,原始地区),肿瘤核心或TC (ET +肿瘤坏死),和整个肿瘤或WT (ET + NT + ED)。网络的最后一层的softmax非线性被sigmoid激活函数所取代，将每个体素作为一个多类分类问题。由于公共和私人排行榜的计算指标是基于这些区域的，这种基于区域的培训可以提高表现。额外的sigmoid输出添加到每个分辨率除了两个最低的水平，应用深度监督和改善梯度传播到早期层。卷积滤波器的数量被初始化为32个，并且分辨率每降低一倍，最大可达320个。

更大网络和组归一化：
  做的第一个修改是，通过将编码器中的过滤器数量加倍，同时在解码器中保持相同的过滤器，非对称地增加了网络的大小。由于训练数据的数量是前一年的四倍，增加网络的容量将有助于它能够建模更大的数据种类。过滤器的最大数量也增加到512个。改进后的网络结构如图2所示。第二个修改是用组归一化代替所有批量归一化。即使使用混合精度训练，3D卷积网络也需要大量的GPU内存，这限制了在训练中可以使用的批量大小。在低批量的情况下，小组标准化比批次标准化工作得更好，之前的竞赛获胜者也采用了这一方法。除非另有说明，组数设置为32。

  图2：结构较大的U-Net与非对称缩放编码器。顶部和底部的数字分别表示通道数量和特征向量的维数。

轴向注意解码器：
  最后添加的是在解码器中使用轴向注意。自注意力或Transformer是一个突破性的想法，允许学习一个输入序列的自适应注意仅仅基于它自己。自我注意机制最初是在自然语言处理文献中构想并推广的，现在已经慢慢被计算机视觉研究团体所采用。当试图将自注意力应用于视觉问题时，主要的障碍之一是注意机制的计算复杂度与输入的大小成二次方，这使得它不可能适合或训练网络在一个标准的工作站设置。当处理带有额外维度的3D数据时，这是一个更大的问题。轴向注意最近被提出作为将注意应用于多维数据时的一种有效解决方案。通过将自注意力独立地应用于输入的每一个轴上，计算只与图像大小成线性比例，使注意机力制即使与3D数据整合也成为可能。作者将轴向注意应用到网络的解码器上，将其运行在转置卷积上采样的输出上，然后将它们相加。
  图3显示了轴向注意解码器块的示意图。即使有了更有效的注意，作者发现这种方法也不可能应用于最高分辨率的特征(128 × 128 × 128)，所以作者只选择了四个较低分辨率的特征。每降低一次分辨率，注意头的数量和每个头的尺寸就增加一倍，分别从4和16 (64×64×64分辨率)开始。

  图3：上采样与轴向注意:轴向注意应用于上采样结果的每一个轴从转置卷积。将输出添加回原始输入，然后将原始输入与编码器路径(未显示)的特性连接起来。

训练：
  遵循nnU-Net针对所有网络的训练方法，每个网络都接受了5倍交叉验证的训练。在训练过程中，动态地应用数据增强来提高泛化能力。数据增强包括随机旋转和缩放、弹性变形、附加亮度增强和伽玛缩放。优化的目标是二进制交叉熵损失和骰子损失的总和，计算在最终的全分辨率输出以及在低分辨率的辅助输出。使用dice损失的批量版本来代替样本dice损失，将整个批次作为一个样本来计算损失，而不是平均每个样本在小批中的dice。批量dice帮助稳定训练通过减少来自样本的错误，少量的注释样本。网络采用Nesterov动量为0.99的随机梯度下降法进行优化。初始学习率为0.01，并按照多项式计划衰减

  每次训练运行持续1000个epoch，每个epoch包含250个小批量。当前fold验证集上的dice得分用于监视训练进度。所有实验都是在24GB VRAM的NVIDIA RTX 3090 GPU上使用Pytorch 1.9进行的。开发了以下模型:

  BL: baseline nnUNet, batch normalization with batch size of 5
  BL + L: baseline with large Unet, batch size of 2, train on all training Samples
  BL + GN: baseline with group normalization, batch size of 2
  BL + AA: baseline with axial attention, batch normalization, batch size
of 2
  BL + L + GN: nnUNet with larger Unet, group normalization, batch size
of 2

实验：

  表1：基于5倍交叉验证的网络dice参数。
  ET:增强肿瘤，TC:肿瘤核心，WT:全肿瘤。

  表1显示了交叉验证5个模型的3个肿瘤子区域的骰子得分。尽管由于在整个训练数据集上进行训练，BL + L模型对于这个比较是无效的，但是增加U-Net的大小产生了微小的改进。使用组归一化而不是批量归一化并没有提高性能，甚至对dice指标造成了轻微的损害。还值得注意的是，使用组归一化会适度增加内存消耗，这可能抵消使用更小的批处理时减少的内存。使用轴向注意编码器并没有提高性能，即使在更高的计算。结合大U-Net和组归一化，对肿瘤核心和整个肿瘤的性能略有提高，GPU内存占用大幅增加。

  表2显示了由竞赛组织者计算的Dice和95% Hausdorff距离(HD95)，并显示在公众排行榜上。除了BL + L配置外，所有预测的分割标签均由5个折叠组合而成。当ET体素的数量小于200时，将增强肿瘤类转化为坏死肿瘤进行后处理。轴向注意模型的结果没有及时为本文处理，故从表中删除。与基线模型相比，合并被提议的变更导致了对所有指标的微小但一致的改进。

  表2：由在线评估平台计算出网络的Dice和95% Hausdorff距离。ET:增强肿瘤，TC:肿瘤核心，WT:全肿瘤，HD95: 95% Hausdorff距离。

定性结果：
  图4显示了BL + L + GN构型的2个代表性预测例子。对于第一个病例，该网络成功地识别了所有的肿瘤子区域，并具有较高的准确性。这可能与扫描质量有关，肿瘤区域对比清晰。在第二个病例中，该网络未能分割增强肿瘤和肿瘤核心，但在整个肿瘤范围内仍表现良好。再次，这可以通过MR扫描的质量来解释:T1和T1- gd对比不能清晰地描绘肿瘤，甚至有可见的伪影和模糊。这部分说明了保证高质量的数据对于网络正常运行的重要性。它还表明，在开发网络时，应该考虑数据完整性因素，无论是通过更多样化的数据采集，还是通过更有意义的数据增强来覆盖这些边缘情况。


  图4：有代表性的好的和坏的预测。T1- gd, T2，和T1，从网络预测两种情况下，好的和坏的指标。绿色:坏死肿瘤核心，蓝色:gd增强肿瘤，黄色:水肿组织。

  在这篇简短的文章中，作者描述了2021年BraTS比赛的方法。通过使用更大的网络扩展nn-UNet框架，用组归一化替换批归一化，并使用轴向注意解码器。这些微小的修改稍微改善了基线nn-UNet。由于nn-UNet框架被广泛用于新的方法，目的是想要突出它的通用性、易用性和健壮性。作者能够在没有大量实验的情况下建立一个非常强的基线。事实上，根据去年的数据开发的预训练网络在今年的验证数据上仍然表现良好，显示了该方法的泛化性。
  对于3D数据，任何修改都需要仔细地平衡GPU内存的可用性，以确保训练能够在没有内存问题的情况下运行。作者提出的更大的U-Net和轴向注意解码器可以为模型增加大量的内存占用，即使进行很小的调整，所以应该仔细调整。组标准化通过允许使用更小的批处理大小，而不会导致显著的性能下降，在一定程度上缓解了这个问题。
  检查失败案例以理解模型的行为是很重要的。大多数分割不准确的情况类似于图4所示的情况，其中一个对比中存在工件或质量问题。其他Dice评分较差的病例是由于后处理方法，有利于切除小增强肿瘤。简而言之，移除小的增强肿瘤可以提高在排行榜上的排名，因为当真实分割标签没有任何增强肿瘤时，Dice和HD95指标采用二分法，作者观察到，使用这种后处理方法，增强肿瘤的Dice评分略有提高，但HD95评分显著下降。作者怀疑这是由于增加了几次扫描小增强肿瘤，这使后处理有害并使HD95恶化。可能需要考虑更复杂的后处理方法来处理这些情况。