生物医学影像自适应全自动深度学习分割网络nnU-net详解

一、引言

U-Net自2015年提出以来，引用数已经超过14000次，且在生物医学领域的图像分割得到了广泛的应用。nnU-net（No New U-Net）之所以叫做“没有新的U-Net”，是因为作者认为“当Pipline设置合理时，基本的U-Net结构是很难被打败的”。基本的U-Net不局限于最初提出的2D U-Net，而是包括基本的2D U-Net、3D U-Net和级联3D U-Net。当然，作者说得比较客气，实际上就是作者认为2015年后对U-Net的一大堆号称SOTA的改进都是对特定数据集的调参，没有本质上的改进。其证据是在常见的49个器官分割任务中（包含19个生物医学数据集，包含CT\MRI\电镜3种影像模态），nnU-net的分割性能均名列前茅，且超过了不少算法专家精心设计的深度学习算法。

二、解决什么问题？

2.1 解决生物医学图像分割问题。对生物医学图像进行分割具有重要意义，是支撑辅助诊疗、手术规划、神经重建、术中协助、肿瘤监测、随访定量分析等生物医学应用的关键技术。70%的国际生物医学竞赛都与图像分割相关（见图1）。

2.2 解决深度学习生物医学图像分割方法的高门槛问题。不同的图像分割任务需要不同的深度学习网络结构（2D、2.5D、3D、级联3D等）、数据扩增（旋转、缩放、Gamma校正、亮度调整等）、超参数（Batch Size、Learning Rate等）、前处理方法（Normalization、Resampling等）、后处理方法（二值连通域的多重处理）的特殊Pipline设置。这些设置常常依赖于人工智能专家的高级知识和经验，并且常常一个小的超参数设置偏差就可能导致AI模型的分割性能大大降低。特别是对于三维生物医学影像分割问题，不同数据集的模态、分辨率大小、是否各向同性等属性差别巨大，难以设置出一个最佳Pipline进行处理。除此之外，不同硬件的显存也会对参数造成巨大影响，且对算法设计者的相关知识要求较高。比如8GB显存一般只支持3D U-Net设置Batch Size为2或3。有研究表明Batch Size较小时不适合进行Batch Normalization，此时应该进行Instance Normalization。如果算法设计者不知道这一点知识，其模型的训练Loss曲线可能起伏很大甚至可能导致无法训练出稳定结果。

2.3 解决生物医学图像分割没有优秀的Benchmark的问题。已经有14000多篇科研文章引用了U-Net，其中有许多对U-Net的改进研究，包括Residual Conv、Dense Conv、Dilated Convs、Attention、Cascade等对U-Net的改进。由于专业门槛、数据集特殊、大量超参数、Pipline设置不一等现实问题，即使是相关专家也难以对这些研究声称的改进的有效性进行客观判断。

图1 典型的生物医学影像数据集的图像与分割标注[1]

三、怎样解决上述问题？

图2 （a）典型的深度学习分割网络模型的处理流程（b）nnU-net的全自动化处理流程[1]

总得来说，nnU-net是全流程自动化的一个分割框架，分为前处理、模型和超参数设置、五折交叉验证的模型训练、后处理等步骤，如图2所示。对于用户来说，只需要提供原始数据和分割label，nnU-net会进行全自动的处理，并最终输出最佳模型。

3.1前处理

3.1.1标准化。对除了CT模态之外的其他模态数据（比如MRI和电镜数据），都使用Z-score进行标准化。也就是说，每张图片都减去自身图片的均值并除以标准差。对CT模态的影像来说，对前景（器官和病灶）进行基于所有数据的前景信息的标准化。背景就是空气等部分。除此之外，还需要进行Clipping，也就是超过99.5%直方图分布外的像素值都统一设置为99.5%那个阈值的灰度值，低于0.5%直方图分布的像素值都统一设置为0.5%那个阈值的灰度值。

3.1.2数据扩增。使用旋转、缩放、添加高斯噪声、进行高斯模糊、亮度调节、对比度调节、Gamma调节、上下和左右镜像等常见的Data Augmentation操作。

3.1.3插值。总体目标是保证从物理意义上来说各个轴上的体素间距是一致的，也就是要做到三维上的各向同性。体素间距越小，分辨率越高。如果某两个轴的分辨率是第三轴的三倍以上，那么这个数据集就被判定为各向异性。对于各向异性数据，主要处理思路是保留高分辨率的两个轴的信息，插值增加低分辨率的那个轴的信息。

3.2自适应生成网络结构

基本的Conv Block的构成为Conv->Instance Norm->Leaky ReLU。每一个encoder和decoder都由2个Conv Block构成。下采样和上采样均使用基础的strided convolutions即可。作者尝试过max pooling和bi\trilinear upsampling的采样方法，发觉与基础方法无实质性差异。当feature map为4*4（*4）时，不再进行下采样。为了尽可能增大感受野，所以patch size的优先级高于batch size。在显存的束缚下，在3D U-Net中，batch size默认为2，以此尽量增加patch size的大小。

3.3模型选择

基础的网络模型包括2D U-Net、3D U-Net和级联3D U-Net。对正常分辨率的医学影像，使用2D U-Net + 3D U-Net。当分辨率过高导致模型的感受野不足时，会影响分割性能，这时就需要使用2D U-Net + 级联3D U-Net。当Patch Size在单轴上小于1/2长时，就将3D U-Net替换为级联3D U-Net。本文的实验结果表明，3D U-Net的分割性能总体来说是最好的。

模型选择的总体思路是进行五折交叉验证，选择最好的模型。有些时候单模型不够好，所以还需要进行Ensemble，以尝试出最佳的输出模型。这个步骤与打比赛时的套路挺一致。

3.4训练

训练的超参数等相关设置大多是固定的。经验表明，如果不使用这些固定参数，会降低分割的精细程度。Epoch设置为1000，training iteration设置为250。Opimizer选择Stochastic Gradient Descent（initial Learning Rate=0.01，nesterov Momentum=0.99）。Learning Rate选择polyLR（多项式曲线下降）。损失函数选择Dice与交叉熵的均值，经验表明，这个设置可以提升分割精度和训练稳定性。

3.5后处理

后处理会使用一个生物医学影像的先验知识：医学影像中常常只有1个主体目标，比如一个人只有1个心脏左心室。在这个先验知识的指导下，后处理会尝试仅保留五折交叉验证后选择的2个模型的重叠的最大连通域，再进行测试。如果测试效果好，就保留此后处理；反之则不保留。

3.6推理

Validation集是被五折交叉验证中对应的训练模型所推理。Test集是被五折交叉验证训练得到的所有模型的最佳模型或Ensemble模型所推理。推理时使用的Patch Size等于 训练时使用的Patch Size。为了避免缝合伪影，使用1/2 Patch Size的间隔进行邻域推理。

四、效果如何？

作者将本方法在19个常见的生物医学影像数据集（见图3）上进行了测试，包括49个器官分割任务，且均取得了名列前茅的分割性能（见图4）。其分割性能超过了不少算法专家精心设计的分割算法。

图3 常见的19个生物医学影像数据集的属性描述[1]

图4 在49个器官分割任务中，nnU-Net均取得了名列前茅的分割性能（红点是nnU-Net，蓝点是其他方法的分割性能）[1]

五、优点

nnU-Net具有根据数据集特性和显存限制进行自适应网络结构生成和全自动训练的优良特点，且不需要调参。在大家都在忙着进行时髦的网络结构、损失函数等创新时，nnU-Net以极大的魄力暗示许多论文作者所宣称的伟大改进可能实际上并没有十分有效。同时，nnU-Net也通过Pipline的丰富集成证实了原定假设“当Pipline设置合理时，基本的U-Net结构是很难被打败的”，未来所有号称有强大泛化能力的分割网络都可以与nnUnet对比一下，以显示自身确实具有有效的创新。

六、不足

6.1 对于非常规的生物医学影像数据集的分割性能是不够的，比如对CREMI比赛中的大脑神经电镜数据进行密集分割，nnU-Net仅取得第6名。这可能是因为nnU-Net没有充分利用神经纤维很长和神经纤维的拓扑结构等先验信息。

6.2 依然有大量固定参数设置。未来研究可以根据数据集特性进行许多参数的自适应预测。

七、启发

7.1做研究不能满足于调用最新的时髦结构对经典工作进行小创新，也应该多想想ResNet、U-Net级别的创新；

7.2 在特定数据集上的特定网络架构改进也是有用的，不需要感到气馁；

7.3 生物医学影像的模态和先验信息多种多样，应该更加充分地考虑数据特点来进行优于nnU-Net的改进与创新。

参考文献

[1] Isensee F, Jäger P F, Kohl S A A, et al. Automated design of deep learning methods for biomedical image segmentation[J]. arXiv preprint arXiv:1904.08128, 2020.

[2] Isensee F, Petersen J, Klein A, et al. nnu-net: Self-adapting framework for u-net-based medical image segmentation[J]. arXiv preprint arXiv:1809.10486, 2018.

[3] Isensee F, Petersen J, Kohl S A A, et al. nnu-net: Breaking the spell on successful medical image segmentation[J]. arXiv preprint arXiv:1904.08128, 2019, 1: 1-8.

[4] Isensee F, Maier-Hein K H. An attempt at beating the 3D U-Net[J]. arXiv preprint arXiv:1908.02182, 2019.

[5] Ronneberger O, Fischer P, Brox T. U-net: Convolutional networks for biomedical image segmentation[C]//International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015: 234-241.

[6] Isensee F, Maier-Hein K H. OR-UNet: an Optimized Robust Residual U-Net for Instrument Segmentation in Endoscopic Images[J]. arXiv preprint arXiv:2004.12668, 2020.

[7] Çiçek Ö, Abdulkadir A, Lienkamp S S, et al. 3D U-Net: learning dense volumetric segmentation from sparse annotation[C]//International conference on medical image computing and computer-assisted intervention. Springer, Cham, 2016: 424-432.