【论文笔记 医疗影像分割—nnUNet】nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation

说明：本文是对原版论文和一位大神解读的基础上，加以自己的理解而作，如有错误，欢迎指正。 大神的文章链接

1.Abstract

对于深度学习模型来说，当用在一个新的问题上，就需要对可变设置进行人为设定。对新问题的适应包括精确架构、预训练、训练、推理对个自由度，这些选择对整体性能有很大的影响。

本文提出nnU-Net（no-new-net），是一种基于三个模型：2D U-Net, 3D U-Net 和U-Net Cascade（级联的3D U-net，之后会有介绍）上的自适应框架。使用者只需要输入数据，nnUNet便可以自动根据图像几何结构调整自身结构，实现全自动分割流程：：

• 预处理：重采样、标准化
• 训练：loss，优化器、数据增强
• 推理：基于patch的策略， ensembling across testtime augmentations and models
• 潜在后处理：enforcing single connected components if applicable

本文对原始U-net仅有微小改动，未使用残差连接、密集连接或注意机制。

2. 网络架构

模型提供了三种Unet架构：2D U-Net, 3D U-Net and a U-Net Cascade(U-Net级联)。三者特点如下：

• 2D/3D Unet输入全分辨率图像；
• Unet Cascade首先用低分辨率图像粗分割，再用全分辨率图像细分割。

2.1 Unet优点（写论文参考）

U-Net[6]是近年来备受关注的一种成功的编解码网络。其编码器部分的工作原理与传统的分类CNN相似，它以减少空间信息为代价，连续地聚集语义信息。由于在分割过程中，语义和空间信息对网络的成功都至关重要，因此必须以某种方式恢复丢失的空间信息。U-Net通过解码器来实现这一点，解码器接收来自“U”底部的语义信息，并将其与通过跳过连接直接从编码器获得的更高分辨率特征图进行重新组合。与其他分割网络（如FCN[9]和以前的DeepLab迭代[10]）不同，这使得U-Net能够很好地分割精细结构。

对原始unet的小改动：①激活函数：用 leaky ReLUs ②归一化方法：使用instance normalization[11]代替了更流行的batch normalization [12]

2.2 2D U-Net

3D医学图像分割使用2D U-Net似乎不是最佳选择，忽略沿z轴的有价值的信息

特点：如果数据集是各向异性的，则常规3D分割方法的性能会下降

2.3 3D U-Net

3D U-Net似乎是选择3D图像数据的适当方法

训练方法：理想情况下，在整个患者的图像上训练；实际上，GPU内存量的限制，只能在切成图像块（image patches）训练。小结构（例如脑肿瘤，海马和前列腺）适合切块，但是像肝这种大结构，切了之后就会损失很多上下文信息。

特点：适合小尺寸物体的分割

2.4 U-Net Cascade

特点：解决3D U-Net在具有大图像尺寸的数据集上的实际缺陷

首先要对降采样的图像进行3D U-Net训练（阶段1）。然后，将阶段1结果上采样到原始尺寸，作为附加输入通道传递到第二个3D U-Net，第二个3D U-Net在上以全分辨率进行训练（第2阶段）。

总结：这里给出了三种算法各自的优缺点，可以针对自己的任务进行选择。也可以都跑一遍，看哪个方法更优。

3. 动态调整网络拓扑

提出原因：由于不同数据集图片尺寸的明显不同，输入patch大小和池化操作数必须自动适应每个数据集，使空间信息的充分聚合。除了适应图像几何形状之外，还有一些技术限制，比如可用的内存需要考虑。

具体实现：根据硬件支持能力进行网络配置。下面以2D Unet为例：

• 2D Unet：
  • 输入patch大小：256x256
  • batch size：42
  • 自动将输入patch大小调整为数据集in-plane大小的中位数
  • 沿轴池化直到这个轴上的特征图大小小于8（但是池数不超过6次）
• 3D U-Net
  • patch size 128 × 128 × 128
  • 如果数据集size中位数<128³，则用数据集size作为patch size，否则128³
  • batch size 2
  • 最高层30个特征图（每次降采样增加2倍）
  • 沿轴池化直到这个轴上的特征图大小小于8（但是池数不超过5次）
• 对于任何网络，我们调整batchsize（最小为2），将每次优化器步骤处理的体素总数（定义为输入补丁体积乘以批处理大小）限制最大为dataset的5%。
• 第1阶段数据集生成的所有网络拓扑;3D U-Net低分辨率是指U-Net级联的第一阶段。3D U-Net Cascade的第二阶段的配置与3D U-Net相同。

根据动态调整网络拓扑策略所获得的网络配置

4. 预处理

预处理是我们的方法组成的全自动分割流程的一部分，下面展示的步骤是在没有任何用户干预的情况下进行的。

4.1 Crop—剪裁

所有数据都裁剪到非零值区域。 这对大多数数据集（例如肝脏CT）没有影响，但会减少颅骨剥脱大脑MRI的大小（从而减少计算负担）。

4.2 Resampling—重采样

CNN不了解体素间距。医学图像中不同协议的slice的采集间隔是不同的。为了使网络更好的学到空间语义信息，将所有患者重新采样到其各自数据集的体素间距的中值，其中三阶样条插值用于图像数据，最近邻插值用于相应的分割蒙版。

U-Net Cascade级联的必要性由下方式决定：如果重采样数据集的shape的中值超过3D U-Net输入patch处理voxels量限度的4倍（批大小为2），它符合U-Net级联，并且该数据集另外被重新采样到更低的分辨率。这是通过增加体素间距（降低分辨率）2倍来实现的，直到达到上述标准。如果数据集是各向异性的，则首先对高分辨率的轴进行降采样，直到它们与低分辨率的轴/轴相匹配，然后才对所有的轴同时进行降采样。

4.3 Normalization

因为CT扫描的强度范围是绝对的，所有CT图像都根据整个各自数据集的统计数据自动归一化：①对训练数据集分割掩码内像素强度的均值和方差进行zscore归一化，并裁剪这些强度值的到[0.5，99.5]。

注：zscore标准化公式为：$z=\frac{x-\mu }{\sigma }$

对于MRI或其他图像模态，则将简单的zscore归一化分别应用于患者。

如果裁剪将数据集中（体素）中患者的平均大小减少了1/4或更多，则只在非零元素的掩码中进行归一化，并且掩码之外的所有值都设置为0。

5. 训练程序

5.1 损失函数

损失函数为diceloss+交叉熵损失函数： $L=L_{dice}+L_{CE}$，其中，dice loss公式如下:

5.2 学习率—当某指标不再变化调整学习率

当训练loss在过去30个时期内没有改善至少5×10−3时，学习率就会降低5因子。如果测试loss在过去60个时代内没有提高超过5×10−3，但在学习率小于10−6之前，训练将自动终止。

5.3 数据增广

在训练过程中，以下增强技术被实时应用（applied on the fly during training）：随机旋转，随机缩放，随机弹性变形，伽玛校正增强和镜像。

如果3D U-Net的输入面片大小的最大边缘长度大于最短长度的两倍，则应用三维数据增强可能不是最佳选择。 对于应用此标准的数据集，我们改用2D数据增强，然后将其逐片应用于每个样本。

U-Net级联的第二阶段接收上一步的分段作为附加输入通道。 为了防止强烈的共适应，我们应用了随机的形态学运算符（腐蚀，扩张，打开，闭合），并随机删除了这些分割的连接部分。

5.4 Patch Sampling

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

6. 推理

推断也都基于patch进行。 由于网络准确性朝着patch的边界降低，因此，当汇总跨patch的预测时，我们将靠近中心的体素的权重提高为接近边界的体素。 通过按patch大小/ 2选择patch以使其重叠，我们通过沿所有有效轴镜像所有patch来进一步利用test time data augmentation(测试时间数据扩充, 也叫TTA, 是一种在比赛中常用的方法)。

对于测试用例，我们使用从训练集交叉验证中获得的五个网络作为整体，以进一步提高模型的鲁棒性。

7. 后处理

使用连通成分分析，删除除了最大连通域之外的其他连通域，减少假阳性区域

8. 结果