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个模型)

后处理

主要就是使用连通域分析

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐