nnUNet 详细解读（一）论文技术要点归纳

nnUNet 论文技术要点归纳

nnUNet

摘要

提出了一种基于2D和3D UNet的自适应框架nnUNet，无需手动调参，平均DICE最高。

介绍

总共10个数据集，7个用于训练阶段，3个用于评估阶段，且不能手动调参。

作者提出的nnUNet，基本没改变UNet的结构，主要工作中心在于：

• 预处理：resampling和normalization
• 训练：loss，optimizer设置、数据增广
• 推理：patch-based策略、test-time-augmentations集成和模型集成等
• 后处理：增强单连通域等

方法

网络结构

• 2D UNet
• 3D UNet
• UNet级联：第一阶段对低分辨率图像进行粗分割，第二阶段进行微调，两个阶段用的都是UNet

修改点

• 激活函数采用leaky Relu
• 归一化采用 instance normalization

网络拓扑动态适应

输入图像大小 不是固定的，为了节省资源，因此需要动态权衡批次大小和网络容量。例如：

3D UNet:

• crop-size<=128x128x128（中值尺寸小于128时，采用中值尺寸
• batch-size>=2
• base_channel=30
• pooling to size>=8
• poolingnum<6

2D UNet：

• crop-size<=256x256（中值尺寸小于256时，采用中值尺寸）
• batch-size<=42
• base-channel=30
• pooling to size>=8
• pooling_num<6

预处理

整体数据Crop：只在非零区域内裁剪，目的是为了减少计算消耗

Resample：数据集中存在不同spacing的数据，默认自动归一化到数据集所有数据spacing的中值spacing。原始数据使用三阶spline插值；Mask使用最邻近插值。

UNet 级联采用特殊的Resample策略：中值尺寸大于显存限制下可处理尺寸的4倍时（batch-size=2），采用级联策略，对数据进行下采样（采样2的倍数，直到满足前面的要求）；如果数据分辨率三个轴方向不相等，先降采样高分辨率轴使得三轴相等，再三轴同时降采样直到满足上述要求。

Normalization：

CT：通过统计整个数据集中mask内像素的HU值范围，clip出[0.05,99.5]百分比范围的HU值范围，然后使用z-score方法进行归一化；

MR：对每个患者数据单独执行z-score归一化。

如果crop导致数据集的平均尺寸减小到1/4甚至更小，则只在mask内执行标准化，mask设置为0.

训练

• 所有模型采用5折交叉验证，重头开始训练，损失函数采用DICE损失和交叉熵损失函数。
• 对于3D UNet和UNet级联的第一阶段，每个样本单独计算DICE loss，然后在一个batch上做平均。
• 对于其他网络，计算整个batch的DICE loss
• Adam优化器，学习率3e-4
• 250个batch/epoch
• 学习率调整策略：计算训练集和验证集的指数移动平均loss，如果训练集的指数移动平均loss在30个epoch内减少不够5e-3，则学习率衰减5倍；
• 训练停止条件：当验证集指数移动平均loss在60个epoch内减少不够5e-3，或者学习率小于1e-6，则停止训练。

数据增广

随机旋转、随机缩放、随机弹性变换、伽马校正、镜像

patch采样：

为了增加网络的稳定性，patch采样的时候会保证一个batch的样本中有超过1/3的像素是前景类的像素。这个很关键，否则你的前景dice会收敛的很慢。

推理

所有的推理都是基于patch的。

patch的边界上精度会有损，因此在对patch重叠处的像素进行fuse时，边界的像素权重低，中心的像素权重高；patch重叠的stride为size/2；使用test-data-augmentation（增广方式：绕各个轴的镜像增广）；使用了5个训练的模型集成进行推理（5个模型是通过5折交叉验证产生的5个模型）

后处理

主要就是使用连通域分析