nnU-Net: Self-Adapting Framework for U-Net-Based Medical Image Segmentation

Github PDF

Abstract

​ U-net在2015年被提出来，到现在引用量已经1万大几了，真是厉害的很。目前在医学图像分割甚至在自然图像领域都应用广泛。但是U-net在针对不同任务的时候，它的网络结构，预处理，训练和推断可能都不同。这些选择不是相互独立的并且影响着最终的效果。作者提出了nnU-Net(no new-Net，起名字就是这么潇洒，学着点)，该网络是在2D和3D U-Net基础的一个鲁棒并且自适应的框架。作者认为不应该过多考虑网络设计的细节，而是应该关注于能提高模型性能和泛化性的其他方面。然后作者进行了实验发现效果惊人的好，我们就好好剖析一下，看看是怎么个惊人法。

Introduction

​ 如果一个网络对指定的任务没有进行完全地手动优化，从而可以有大量的空间进行调整去改善结果，那么很容易证实这些用来提高网络性能的结构调整可以很有效（翻译的有点绕口）。在作者之前的实验中，这些调整并没有提高分割的性能并且不能推进SOTA（感同身受啊，我之前做胰腺的CT 3D分割，无论怎么加模块（记得当时用了各种很火的注意力机制，并且看很多新论文又推出了什么什么牛B的模块），调整结构，效果怎么也不能再往上提高了，一度怀疑我是否已经达到了SOTA )。这让我们相信非结构因素在分割方法中更重要，同时它也被严重低估（记得我当时也用了很多预处理方法，还有在线的数据增强）。nnU-Net使用了三个相对简单的U-Net，并且这些U-Net只是对原始U-Net进行了很小的修改，并且没有使用各种扩展插件（残差连接、Dense连接还有各种注意力机制）。nnU-Net可以自动将其架构适应给定的图像几何，更重要的是，它彻底定义了围绕他们的所有其他步骤（听着很霸气）。这些步骤包括：

1. 预处理，比如resampling和normalization
2. 训练，比如损失函数，优化器的设置和数据扩充
3. 推断，比如基于图像块的策略，TTA（test-time augmentation）集成和模型集成
4. 后处理，比如增强单连通域

Methods

网络结构

​ 作者提出了3个网络，分别是2D U-Net、3D U-Net和级联Unet。2D和3D U-Net可以生成全分辨率的结果，级联网络的第一级产生一个低分辨率结果，第二级对它进行优化。

• 改动：使用Leaky ReLU(neg.slope 1e-2)取代ReLU，使用Instance normalization取代Batchnormalization(解释)
• 2D U-Net: 既然使用了3D U-Net，为啥还有用2D的。因为有证据（这证据也是作者发现的）表明当数据是各向异性的时候，传统的3D分割方法就会很差（怪不得我之前的效果一直上不去，难道是因为我用3D分割方法的原因？）
• 3D U-Net: 3D网络固然好，就是太占用GPU显存。那我们可以使用小一点的图像块去训练，但这对于那些较大的图像例如肝脏，这种基于块的方法就会阻碍训练。这是因为受限于感受野的大小，网络结构不能收集足够的上下文信息去正确的识别肝脏和其他器官。
• 级联3D U-Net：为了解决3D U-Net在大图像尺寸数据集上的缺陷，本文提出了级联模型。首先第一级3D U-Net在下采样的图像上进行训练，然后将结果上采样到原始的体素spacing。将上采样的结果作为一个额外的输入通道（one-hot编码）送入第二级3D U-Net，并使用基于图像块的策略在全分辨率的图像上进行训练。

网络拓扑的动态自适应

​ 由于输入图像大小的不同，输入图像块大小和每个轴池化操作的数量（同样也是卷积层的数量）必须能够自适应每个数据集去考虑充足的空间信息聚合。除了自适应图像几何，还需要考虑显存的使用。指导方针是动态平衡batch-size和网络容量。

​ 首先展示一下网络初始配置：

• 2D U-Net, 图像大小=256X256，batch size=42，最高层的特征图谱数量=30（每个下采样特征图谱数量翻倍），自动将这些参数调整为每个数据集的中值平面大小（这里使用面内间距最小的平面，对应于最高的分辨率），以便网络有效地训练整个切片。我们将网络配置为沿每个轴池化，直到该轴的特征图谱小于8（但最多不超过6个池化操作）。
• 3D U-Net, 图像大小=128X128X128, batch size=2, 最高层的特征图谱数量=30。由于显存限制，不去增加图像大小超过128³体素，而是匹配输入图像和数据集中体素中值大小的比率。如果数据集的形状中值比128³小，那就使用形状的中值作为输入的图像大小并且增加batch size（目的是将体素的数量和128X128X128大小，batch size为2的体素数量相等）。沿每个轴最多池化5次直到特征图谱大小为8。

在每个优化步骤中，我们限制每个网络的体素数量（输入图像的体素 X batch size）最多为数据集的5%。如果超过，则降低batch size(最小为2)。

​ 可以看到不同数据集的输入图像块大小是不同的，3D U-Net lowres 为级联3D U-Net的第一级，第二级的配置和3D U-Net网络一致。

预处理

• Cropping：将所有数据的非零区域进行裁剪

• Resampling：我们知道医学影像数据不同的设备和设置会导致具有不同体素间距的数据，而CNN并不能理解这种体素间距的概念（别说CNN了，我一开始也不懂）。为了让网络学会空间语义，所有的病例重采样到相应数据集的体素间距中值，对数据和mask分别使用三阶样条差值和最近邻差值方法。

  注：如果重采样数据的形状中值可以作为3D U-Net中的输入图像（batch size=2）的4倍以上，则使用级联U-Net网络，且数据集需要重新采样到较低的分辨率。可以以2为因子增加体素间距（降低精度）直到4倍那个条件不满足。如果该数据集是各向异性的（比如1.2mm X 1.2mm X 3.5mm），首先更高分辨率的轴进行下采样直到与低分辨率轴匹配，然后才对所有轴同时进行下采样。

• Normalization：对于CT图像，首先搜集分割mask内的像素值，然后所有的数据截断到这些像素值的[0.5, 99.5]%，然后进行z-score标准化；对于MRI图像，直接进行z-score标准化。

  注：如果因为裁剪减少了病例平均大小的1/4或更多，则标准化只在非零元素的mask内部进行，并且mask外的所有值设为0。

训练

• 5折交叉验证

• Loss = Dice + Cross Entropy

  

  u为模型概率输出，v是ground Truth的one hot编码，k为类别数。

• 优化器：Adam，初始学习率3e-4，每个epoch有250个batch。

• 学习率调整策略：计算训练集和验证集的指数滑动平均loss，如果训练集的指数滑动平均loss在近30个epoch内减少不够5e-3，则学习率衰减5倍。

• 训练停止条件：当学习率大于10^-6且验证集的指数滑动平均loss在近60个epoch内减少不到5e-3，则终止训练。

• 数据扩充（On the fly）：随机旋转，缩放，elastic deformation，gamma矫正和镜像。代码

  注：如果3D U-Net的输入图像块尺寸的最大边长是最短边长的两倍以上，这种情况对每个2维面做数据增广。

• 级联U-Net的第二级接受前一级的输出作为输入的一部分，为了防止强co-adaptation，应用随机形态学操作（腐蚀、膨胀、开运算、闭运算）去随机移除掉这些分割结果的连通域。

• 图像块采样：为了增强网络训练的稳定性，强制每个batch中超过1/3的样本包含至少一个随机选择的前景。

推断

​ 所有的推断都是基于图像块（patch-wise）。网络精度沿着图像块边缘逐渐降低，所以提高图像块中心的体素权重，降低边缘权重。图像块重叠大小为size / 2，并使用TTA（沿各个轴镜像）进行集成。在测试集上测试时使用5个交叉训练的模型进行集成来提高模型的鲁棒性。

后处理

​ 一般认为某一类是在单连通域内，也就是说一个病例内只会有一个这样的区域。那么只保留最大的连通域，将其他的所有小连通域去掉。

实验

竞赛官网

​ 这个竞赛的目的就是开发一套方法能适应不同的医学影像数据，下面是作者在训练集上交叉验证的结果和在测试集上的结果。

​ 可以看到除了在脑肿瘤分割任务中测试集上性能下降有点大，在其他数据集上简直就是完美（相对于在训练集的交叉验证结果）。作者也解释了这是因为脑肿瘤的数据和Ground Truth之间存在差异（怪数据，不怪模型）。在竞赛官网可以看到本文排名第一（哄哄）。

总结

这里借用张良怀的总结和分析

​ 作者也说了，文章中的各种改进也并没有理论实验支撑，比如Leaky ReLU替换ReLU，数据扩充方法。因此需要做更多的消融实验去证明。

作者有更详细的版本Automated Design of Deep Learning Methods for Biomedical Image Segmentation，推荐精读。

---

我的感悟

​ 通读这篇文章的原因是因为他的成绩很好，想一探究竟。回想我们以前去做这样的任务的时候，往往会在模型上花更多的时间。比如添加各种模块，像常用的残差连接，Dense连接，还有各种各样的注意力机制，或者增加减少特征图谱的数量。有时候这种模型的改进似乎真的提高了效果，但归根结底这种效果也只是拼凑出来的，具体为什么能这样，论文作者也说不出个所以然来，而且模型的泛化性值得考量。本文作者就是用最原始的U-Net通过模型调整以外的一些操作来达到更好的效果。当然这有点打比赛的意味，为了最终结果不考虑测试时间，离最终的临床似乎有点偏离。但我觉得这是一种进步，一种跳出思维定式的进步。他打碎了我们对网络结构调整的这种痴念，让我们回归到正常的科研道路上。当然不是说网络结构调整没有用，只是我觉得现阶段网络结构调整应该是一种进化式的，比如残差网络。而说不好听一点目前很多网络的改进只是为了灌水（当然作者本意可能不是如此），没有太多实际价值。