提出了一种基于2D和3D UNet的自适应框架nnUNet,无需手动调参,平均DICE最高。
总共10个数据集,7个用于训练阶段,3个用于评估阶段,且不能手动调参。
作者提出的nnUNet,基本没改变UNet的结构,主要工作中心在于:
- 预处理:resampling和normalization
- 训练:loss,optimizer设置、数据增广
- 推理:patch-based策略、test-time-augmentations集成和模型集成等
- 后处理:增强单连通域等
输入图像大小 不是固定的,为了节省资源,因此需要动态权衡批次大小和网络容量。例如:
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采样的时候会保证一个batch的样本中有超过1/3的像素是前景类的像素。这个很关键,否则你的前景dice会收敛的很慢。
所有的推理都是基于patch的。
patch的边界上精度会有损,因此在对patch重叠处的像素进行fuse时,边界的像素权重低,中心的像素权重高;patch重叠的stride为size/2;使用test-data-augmentation(增广方式:绕各个轴的镜像增广);使用了5个训练的模型集成进行推理(5个模型是通过5折交叉验证产生的5个模型)
主要就是使用连通域分析