【论文笔记】nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation

nnU-Net:基于U-Net的自适应医学图像分割框架

Background

先前的医学图像分割模型都需要专门的架构和训练方案来实现最优的性能，通常只在少数甚至单个数据集上能够有较好的分割效果。

Novelty

使用全自动的方式调整网络架构和其他步骤（如预处理、数据增强、后处理等）以适应新的分割任务，无需手动交互。

训练时使用五折交叉，将生成的五个网络合并为一个整体进行最终的预测，可以提高模型的鲁棒性。

Contribution

1. 基于U-Net提出了nnU-Net框架，可以自动调整模型的体系结构以适应不同数据集。

2. 作者认为现有的模型在结构方面的调整可能会导致模型在某一些特定的数据集上过拟合，同时非架构的部分在图像分割任务中更重要。在常见的生物医学影像数据集上进行了测试，仅对原始U-Net进行微小调整的nnU-Net均取得了良好的分割性能，超过了很多修改的架构。

Method

网络架构

总体结构包含三个网络：2D U-Net、3D U-Net、U-Net Cascade。2D U-Net和3D U-Net生成全分辨率分割，U-Net Cascade先生成低分辨率分割，然后对其进行细化。在交叉验证后，nnU-Net根据经验选择最佳性能配置或集成。

使用leaky ReLUs代替ReLU，instance normalization代替batch normalization。

U-Net Cascade的结构如下：

先将数据下采样，经过3D U-Net得到低分辨率的分割图，再上采样到原始分辨率，与全分辨率图像连接，经过第二个3D U-Net得到全分辨率的分割图。

网络拓扑结构的动态适应

由于图像大小的巨大差异，patch size和每个轴的池化操作数、卷积层数必须自适应每个数据集，以允许充分聚合空间信息。

首先将patch size设定为原始图像大小的中位数，然后迭代的减少patch size，同时调整网络的拓扑结构（网络深度、池化操作数量、池化操作位置、feature map 的尺寸、卷积核尺寸等），直到网络在显存限制下可以被训练，且batch size不小于2。

预处理

裁剪：所有图像剪裁到非零值区域。

重采样：将所有图像重采样至数据集体素间距的中位数，图像数据使用三次样条插值，mask使用最近邻插值。

应使用U-Net cascade模型的数据：心脏、肝脏、肺和胰脏。

正则化：对于CT扫描图，训练集中所有segmentation mask中的value会被收集，整体的数据集会先被clip到[0.5%,99.5%]，然后通过收集的数据的均值和标准差进行正则化。

对于MRI以及其他的图像，直接进行正则化。

训练过程

在训练集上使用五折交叉验证进行评估，loss使用Dice和CE的组合。

Dice使用的是适应多类的变体：

 

（u为soft-max输出，v为ground truth的one-hot编码，I为训练patch/batch中的像素个数，K为类）

使用Adam优化器初始learning rate为0.0003，将超过250个batch的迭代定为epoch，当训练集平均loss在30个epoch内下降低于0.005时学习率降低五倍，验证集的loss在60个epoch内降低未超过0.005且学习率小于10^-6，则终止训练。

数据增强

使用随机旋转、随机缩放、随机弹性变形、伽马校正增强和镜像。

如果三维U-Net的输入patch size的最大边长度是最短边长度的两倍以上，则应用三维数据增强可能是次优的。应对样本的每个切片使用2D数据增强。

U-Net Cascade 的第二阶段接收上一步的分割图作为附加输入通道。 为了防止强协同适应，应用随机形态算子（腐蚀、扩张、打开、关闭）并随机删除这些分割图的连接分量。

patch 采样

为了增加网络训练的稳定性，强制在一个batch中超过三分之一的样本包含至少一个随机选择的前景类。

推理

patch的边界网络准确性较低，因此，当汇总跨patch的预测时，需要将靠近中心的体素的权重提高。通过按patch size/ 2选择patch以使其重叠，通过沿所有有效轴镜像所有patch来进一步利用test time data augmentation(测试时间数据扩充)。

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

后处理

通过连通成分分析，删除除了最大连通域的其他连通域，以减少假阳性区域。

不足

训练三个模型并为每个数据集独立选择最佳模型并不是最优的解决方案。给定一个较大的时间尺度，可以在训练之前研究适当的启发式方法，以确定给定数据集的最佳模型。