本文是 nnU-Net 论文的笔记,稍微做了整合了一下,可能有些地方有所遗漏,等以后看完代码之后再补充。
1. 摘要
nnU-Net(no-new-Net)是基于 U-Net 的一种模型,它在达到很好的效果的同时还实现了自适应(self-adapting)的功能。在不同数据集(或不同的部位)的医学图像上进行分割时,往往需要具有不同结构的网络和不同的训练方案,自适应是指模型在对不同的数据集进行训练时,可以自动的调整 batch-size、patch-size 等,以达到很好的效果。
作者在一个医学图像分割十项全能比赛当中的 6 个数据集集中都取到了当时最好的结果,这个比赛一共 10 个数据集,它会给出 7 个数据集让你训练,然后在其他 3 个数据集上进行验证。
2. 方法
2.1 网络结构
作者以 U-Net 为基础,通过调整网络的结构后效果没有得到提升,所以作者转而关注非结构的部分。nnU-Net 可以根据给出的数据集自动的调整网络的设置。nnU-Net 相比于 U-Net,前者用 leaky ReLU 代替了 ReLU 做激活函数,用实例正则化(instance normalization)代替了批正则化(batch normalization)。
U-Net 采用编码器-解码器的结构,其中编码器部分可以聚合语义信息(semantic information),但是同时也会损失空间信息(spatial information),说人话就是编码器在提取图像特征的同时也使图像的分辨率降低了。但是语义信息和空间信息对于分割任务来说都是至关重要的,所以在解码器部分需要对空间信息进行恢复,具体的就是通过下层的语义信息,并结合利用跳跃连接得到的更高分辨率的特征图(即空间信息)来进行空间信息的恢复。
nnU-Net 由三种基础的 U-Net 网络组成,分别是 2D U-Net,3D U-Net 和 U-Net Cascade。其中,2D 和 3D U-Net 产生一个全像素的分割结果,U-Net Cascade 先产生一个低像素的分割结果,再对其进行微调。
直观来看,2D U-Net 在处理 3D 的医学图像时不是最优的,因为在 $z$ 轴上的信息没有得到充分的利用。但是也有证明显示 3D 分割方法在一些具有各向异性(不懂何为各向异性)的数据集上效果会变差。我们通常会选择 3D U-Net 来处理 3D 的医学图像,但这 3D 分割方法的一个缺点就是占用显存较大,当显存放不开一个 3D 图像时,可以将图像的 patch 作为输入,但这又限制了视野(field of view),不能充分的利用图像的上下文信息。为了解决 3D U-Net 带来的这种问题,所以使用了 U-Net Cascade 的方法,即先用一个 3D U-Net 处理下采样之后的图像(图像分辨率较低),然后输出分割结果,并将其上采样到原始大小,并作为附加的输入通道(和原始 3D 图像在通道上拼接)输入到第二个 3D U-Net 做基于 patch 的分割。
上图是 U-Net Cascade 的示意图,要注意每个不同标记所代表的意思,然后结合上一段的内容还是比较好理解的。
2.2 动态自适应
网络的动态自适应简单来说就是根据数据集的不同自动的调整 batch size 和 patch size,当 patch size 大的时候 batch size 就小一点;当前者小的时候,后者就大一点。此外还要调整池化和卷积操作的个数。
3. 预处理
3.1 裁剪
将图像中非零的区域裁掉,可以减小显存的负担。
3.2 重采样
由于成像设备和成像协议的不同,医学图像的体像素间距也不同,所以要对其进行重采样。具体的,对医学图像使用三阶样条插值法,对其相应的分割标签使用最近邻插值法。
3.3 正则化
对于 CT 图像而言,对其像素值归一化到原像素值的 0.5%~95.5%,并根据均值和方差对其进行 z-score 正则化。对于非 CT 图像而言,只进行 z-score 正则化。
4. 训练
4.1 网络设置
训练时的总损失 $= dice$ 损失 + 交叉熵损失,优化器使用的是 Adam 优化器,学习率为 $3\times 10^{-4}$,每当训练损失的指数移动平均值在最近 30 个 epoch 内没有改善时,学习率下降 0.2 倍。当在 60 个 epoch 内,学习率的变化不超过 $5\times10^{-3}$ 并且学习率不小于 $10^{-6}$ 时,停止训练。在验证时使用五折交叉验证。
4.2 数据增强
使用随机旋转、随机缩放、随机弹性变形、伽玛校正增强和反射做数据增强。在第二阶段的 U-Net Cascade 中使用了形态学操作和随机移除分割中的连接组件来进行数据增强。
5. 推断
因为网络的准确率在一个 patch 中的中央部分较高,而在边缘部分较低,所以在在中央部分的体素的权重较高,而边缘部分的较低。
6. 后处理
对 3D 医学图像进行分割,最终器官的分割结果在空间当中一定是连通的,所以最后只保存最大的连通分量,其他的均删除。
- 本文作者: 俎志昂
- 本文链接: zuzhiang.cn/2020/02/21/nnU-Net/
- 版权声明: 本博客所有文章除特别声明外,均采用 Apache License 2.0 许可协议。转载请注明出处!