nnU-Net论文笔记

论文标题：nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation 一种基于U-Net的医学图像分割自适应框架

论文链接：arxiv

摘要（abstract）

提出当Unet处理新问题时需要有适应过程，四个方面：确切的架构，预处理，训练和推理。这些选择并不是相互独立的，会对模型性能产生重大的影响。

The adaptation of the U-Net to novel problems, however, comprises several degrees of freedom regarding the exact architecture, pre-processing, training and inference.

为此，本文提出了nnU-Net(“no-new-Net”) ，一种基于2D或3D Unet的健壮的，自适应的网络。

The present paper introduces the nnU-Net (”no-new-Net”), which refers to a robust and self-adapting framework on the basis of 2D and 3D vanilla U-Nets.

nnU-net认为去除网络设计方案中华而不实的部分，关注能使模型具有良好性能和泛化性的部分是很有必要的。作者在”Medical Segmentation Decathlon“挑战赛中测试了nnU-net的性能，在所有分割类别中达到了最高的平均Dice分数。

介绍（Introduction）

首先提出当今绝大多数模型只对某些特定数据集具有良好表现，这使得研究人员如果不能很好地处理模型的细节就无法得到论文所提到的优良性能。接着介绍“Medical Segmentation Decathlon”挑战赛是一个包含了不同数据集和数据实体用来避免模型在特定数据集上具有良好表现，参赛模型只被允许进行全自动的自适应特定数据集，不能人工干预的比赛。


挑战分为两个连续的阶段：1) 开发阶段，参与者可以访问 7 个数据集以优化他们的方法，并且使用他们最终的冻结方法，必须为相应的 7 个保留测试集提交分段. 2) 第二阶段在 3 个以前未公开的数据集上评估相同的已训练好的模型。

作者提出假设（hypothesize）：最近提出的一些架构的修改模型会导致模型针对特定数据集的过拟合或者在最优表现的基础上次优化导致验证集效果不完美。

We hypothesize that some of the architectural modifications presented recently are in part overfitted to specific problems or could suffer from imperfect validation that results from sub-optimal reimplementations of the state-of-the art

作者认为，在还没达到最佳优化的架构中通过微调是可以实现模型性能的提升，但是在作者自己的初步实验中发现，对已达到最优化的架构中进行调整并不能提升模型的性能。因此，作者相信非架构方面对分割方法的影响大得多，同时该方面也被严重低估了。
至此作者提出了nnU-net，它建立在一组三个相对简单的U-Net模型上，这些模型仅仅对原始的unet进行了小小的修改。nnU-Net 自动调整其架构以适应给定的图像几何。

方法（Methods）

1. 网络架构

医学分割图像通常包含三个维度，所以作者采用2D U-Net、3D U-Net和U-Net cascade组成基本的nn U-net架构池。
U-Net 是一个成功的编码器-解码器网络，近年来受到了广泛关注。它的编码器部分与传统分类 CNN 的工作方式类似，因为它以减少空间信息为代价连续聚合语义信息。由于在分割中，语义和空间信息对于网络的成功都是至关重要的，因此必须以某种方式恢复丢失的空间信息。 U-Net 通过解码器执行此操作，解码器从“U”的底部接收语义信息，并将其与通过跳过连接直接从编码器获得的更高分辨率的特征图重新组合。与其他分割网络不同，例如 FCN 和 DeepLab 的先前迭代，这允许 U-Net 特别好地分割精细结构。

与原U-Net不同的是，作者用 leaky ReLU（负斜率 1e−2）替换 ReLU 激活函数，并使用实例归一化 而不是更流行的批量归一化 。

为了解决 3D U-Net 在具有大图像尺寸的数据集上的视野受限，不能充分利用上下文信息，我们另外提出了一个（U-Net cascade）级联模型。因此，首先在下采样图像上训练 3D U-Net（第 1 阶段）。然后将此 U-Net 的分割结果上采样到原始体素间距，并作为附加（独热编码）输入通道传递给第二个 3D U-Net，后者在全分辨率的补丁上进行训练（第 2 阶段）。如下图所示。

2.网络拓扑结构的动态调整（dynamic adaption of network topologies）

由于图像大小的巨大差异（肝脏的中值形状为 482 × 512 × 512 与海马体的中值形状为 36 × 50 × 35），输入块大小和每个轴的池化操作数（以及因此隐含的卷积层数）必须自动适应每个数据集，以允许充分聚合空间信息。除了适应图像几何形状之外，还有一些技术限制，例如可用内存。作者在这方面的指导原则是动态权衡批量大小与网络容量，详情如下：

对于 2D U-Net，此配置是输入块大小为 256×256，批量大小为 42，最高层有 30 个特征图（特征图的数量随着每次下采样而加倍）。我们自动将这些参数调整为每个数据集的中值平面大小（我们使用平面内间距最小的平面，对应于最高分辨率），以便网络有效地训练整个切片。对于3D U-Net，此配置的输入块大小为128128128，批量大小为2，最高层有30个特征图。

为第 1 阶段数据集生成的所有网络拓扑如下表所示。

3.预处理
作者提出的预处理是模型的必要部分，但是不同的是，这不是人为干预的。

• cropping：所有数据都裁剪到非零值区域。 这对大多数数据集（如肝脏 ）没有影响，但将减少颅骨剥离脑MRI的尺寸（以及因此产生的计算负担）。
• Resampling：CNN本身并不理解体素间距（voxel spacings）。为了使我们的网络能够正确学习空间语义，所有患者都被重新采样到各自数据集的中位体素间距，其中三阶样条内插（third order spline interpolation）用于图像数据，最近邻内插用于对应的分割掩码。

“像素间距”（Pexel Spacing）是指两个像素之间的距离，spacing定义了图像像素的物理大小并且保证了实际距离测量的准确性。比如，如果知道x和y轴的像素间距为 0.4mm，那么在图像中的一条 10 像素的线就会有 4mm的长度。同样，由于知道图像像素中的宽和高(比如对于普通CT来说是 512×512) ，就能够找到图像的实际尺寸了:512 × 0.4 mm = 204.8 mm。对应的，当图像从2D扩展到3D时， “像素”会扩展成“体素”，“像素间距”（Pexel Spacing）也会扩展成“体素间距” （Voxel Spacing）。综上，图像分辨率一般是指“像素/体素数”，而spacing是指“体素间距”，知道了体素数和体素间距，就知道了一张CT的“尺寸”。

CT图像的重采样
与自然图像不同，在医学影像中，人体部位真实的大小（成像大小）是非常重要的信息。因此例如在CT图像中，存在体素间距（spacing）和体素个数（resolution）两个指标： 成像大小 = spacing * resolution
且成像大小保持不变。
由于不同的扫描仪或不同的采集协议通常会产生具有不同体素间距的数据集，而CNN无法理解体素间距，因此我们需要将所有医学影像的spacing重采样到一致，这样resolution就可以反映成像大小了。
增大CT图像的spacing，其总体的resolution会减少，因此一个patch中获取的细节信息减少，上下文信息增大；而减少spacing，其总体的resolution会增大，一个patch中获取的细节信息会增大，而上下文信息又会减少。因此如何做好这个trade-off在医学影像中是非常重要的。

• Normalization：因为CT扫描的强度范围是绝对的,CT图像都是基于统计的自动归一化各自的数据集。如果裁剪（cropping）将数据集（体素）中的患者的平均规模减少1/4或更多，则仅在非零元素的面具内进行规范化，并且面具外的所有值都设置为0。

训练

作者在训练集上使用“five-fold cross-validation”进行评估，并且使用Dice和交叉熵损失联合作为损失函数：

为了提高网络训练的稳定性，我们强制执行批次中超过三分之一的样本包含至少一个随机选择的前景类。

推论（Inference）

Due to the patch-based nature of our training, all inference is done patch-based as well. Since network accuracy decreases towards the border of patches, we weigh voxels close to the center higher than those close to the border, when aggregating predictions across patches. Patches are chosen to overlap by patch size / 2, and we further make use of test time data augmentation by mirroring all patches along all valid axes.

Combining the tiled prediction and test time data augmentation result in segmentations where the decision for each voxel is obtained by aggregating up to 64 predictions (in the center of a patient using 3D U-Net). For the test cases we use the five networks obtained from our training set cross-validation as an ensemble to further increase the robustness of our models.

实验结果

第一阶段交叉验证结果以及相应的提交测试结果总结，表中“-”表明U-Net级联是不适用的数据集,因为它已经完全覆盖的输入块大小3D U-Net。粗体表示最终作者提交到比赛中的模型精度。

总结

本文介绍了直接围绕原始U-Net架构[6]的医学领域nU-Net分割框架，并动态地适应任何给定数据集的细节。 基于我们的假设，即非架构修改可能比最近提出的一些架构修改更强大，这个框架的本质是自适应预处理，训练方案和推理的彻底设计。 适应新细分任务所需的所有设计选择都是以完全自动的方式完成的，无需手动交互。 对于每项任务，nnU-Net自动为三种不同的自动配置U-Net模型进行five-fold交叉验证，并选择最高平均前景Dice分数的模型（或结合）。