nnUNet论文阅读

nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation

摘要：

The key design choices in this process are modeled as a set of fixed parameters, interdependent rules and empirical decisions. Without manual intervention, nnU-Net surpasses most existing approaches, including highly specialized solutions on 23 public datasets used in international biomedical segmentation competitions. We make nnU-Net publicly available as an out-of-the-box tool, rendering state-of-the-art segmentation accessible to a broad audience by requiring neither expert knowledge nor computing resources beyond standard network training.

设计了一套能够适用于大部分医学图像的处理流程，在不进行人为干涉的情况下，能够超过大部分算法。

In this work, we outline a new path between the status quo of primarily expert-driven method configuration in biomedical segmentation on one side and primarily data-driven AutoML approaches on the other. Specifically, we define a recipe hereafter that systematizes the configuration process on a task-agnostic level and drastically reduces the search space for empirical design choices when given a new task.

1. Collect design decisions that do not require adaptation between datasets and identify a robust common configuration (‘fixed parameters’).
2. For as many of the remaining decisions as possible, formulate explicit dependencies between specific dataset properties (‘dataset fingerprint’) and design choices (‘pipeline fingerprint’) in the form of heuristic rules to allow for almost-instant adaptation on application (‘rule-based parameters’).
3. Learn only the remaining decisions empirically from the data (‘empirical parameters’).

方法：

dataset fingerprints:

1、去除全为0的背景区域

2、根据在crop 前后的image size，image spacing（体素的实际大小）、模态、类别数目、训练数据量创建“dataset fingerprint”，同时 dataset fingerprint 包括前景区域体素的统计信息

pipeline fingerprints：

启发式规则，这些规则包含了领域知识，dataset fingerprints和硬件约束进行操作。这些基于规则的参数由固定参数(与数据无关)和经验参数(在训练过程中进行优化)补充

fixed parameters：

1、网络拓扑结构：采用全原始的UNet或者3D Unet结构，为了更大的patch size，模型的batch size很小，只能为2，使用IN代替BN；leaky ReLU代替ReLU；使用深监督机制；使用跨步卷积进行下采样使用转置卷积进行上采样；初始的通道数为32；同时模型的最大通道数小于320（3D）、512（2D）

训练策略：

训练1000个epoch，每个epoch包括250个mini batch；SGD+动量（0.99）作为优化器；使用‘poly’ learning rate policy32；使用CE loss + DICE loss 作为损失函数；深监督时将GT进行下采样，对不同分辨率的损失加权，分辨率越低，权重越小，指数衰减；为了处理类别不均衡，使用过采样，66.7%的样本来自所选训练样本的中的随机位置，而33.3%的patch保证包含在所选训练样本中存在的前景类之一(随机选取)。前景补丁的数量被四舍五入，强制最小值为1(结果是一个随机补丁和一个前景补丁，批大小为2)。训练时的数据增强在表格中。

推理：

使用滑窗，滑窗大小为训练的patch size，步长为0.5 patch size；高斯加权抑制边缘错误； 每个轴使用镜像增强

Rule-based parameters

1、灰度归一化： 有两种灰度归一化方法，除了CT图像外，所有模式的默认设置都是z-scoring。这个选项,在培训和推理,每个图像归一化独立首先减去它的均值,然后除以std。如果cropping导致平均大小减少25%或更多,创建一个中央非零像素点的mask，对mask使用z-scoring,忽略了周围的零体素。对于CT图像，nnU-Net采用了不同的方案，因为强度值是定量的，反映了组织的物理性质。因此，通过使用适用于所有图像的全局归一化方案来保留这些信息是有益的。为此，nnU-Net使用前景体素的0.5和99.5个百分点进行裁剪，同时使用全部前景均值和s.t.d.对所有图像进行归一化。

2、重采样：

target spacing是一个关键参数。较大的spacing导致较小的图像，从而丢失细节，而较小的间隔导致较大的图像，阻止网络积累足够的上下文信息，因为 patch size受到GPU显存的限制。虽然3D U-Net级联部分解决了这个问题，但对于低分辨率和全分辨率，仍然需要合理的目标间距。对于3D全分辨率的U-Net, nnU-Net使用每个轴独立计算的训练数据中找到的spacing中位数作为默认target spacng。对于各向异性数据集，这种默认值可能会导致严重的插值伪影，或者由于训练数据的分辨率差异很大而导致大量信息丢失。因此，如果体素和spacing各向异性(即最低间距轴与最高间距轴的比值)都大于3，则选择最低分辨率轴的target spacing为训练数据的 spacing的十分位数。对于2D U-Net，一般采用nnU-Net在两个轴上操作，分辨率最高。如果三个轴都是各向同性的，则两个尾部轴用于切片提取。target spacing是训练数据的中位数(对每个轴独立计算)。对于基于切片的处理，不需要沿面外轴进行重采样。

3、Adaptation of network topology, patch size and batch size

大的patch size，batch size 最低为2

同时考虑图像spacing，可以旨在指定轴上进行降采样或者。所有U-Net配置的网络拓扑都是根据重采样后图像的中值大小以及重采样图像的target spacing来选择的。适应化过程流程图见Supplementary Note 5中的图SN5.1。由于GPU内存占用估计是基于feature map大小进行的，因此适配进程不需要使用GPU。

4、初始化

patch size需要被整除，如果不被整除，则需要pad

5、网络结构

每个轴下采样操作数目取决于patch size与target spacing。

向下采样直到进一步向下采样将feature map的大小减少到小于4个体素，或者feature map的间距变得各向异性。下采样策略由target spacing决定;高分辨率轴分别向下采样，直到它们的分辨率在低分辨率轴的两个因子内。随后，所有轴同时向下采样。每个轴分别终止下采样直至触发各自的特征图约束。对于3D U-Net和2D U-Net，卷积的默认内核大小分别为3 × 3 × 3和3 × 3。如果在轴之间存在初始分辨率差异(定义为大于2的间距比)，则平面外轴的核大小设置为1，直到分辨率在2倍以内。注意，对于所有轴，卷积核的大小都保持为3。

6、使用GPU显存

配置过程中可能的最大patch size受GPU内存的限制。由于重新采样后，patch size被初始化为图像的中值形状，所以对于大多数数据集来说，它最初太大了，无法适合于GPU。nnU-Net根据网络中特征图的大小估计给定架构的内存消耗，并将其与已知内存消耗的参考值进行比较。patch size就会减少在一个迭代的过程,而架构配置相应的更新在每一步,直到达到所需的预算(图SN5.1补充注5)。减少patch size总是应用于最大轴相对于中值图像数据的形状。一步的减少等于该轴的第二个体素，其中nd是下采样操作的数量。

7、Batch size

如果执行了减小patch size的操作，则 batch size 设置为2。否则，将使用剩余的GPU内存空间来增加批大小，直到GPU被完全使用。为了防止过拟合，批大小是有上限的，这样小批中的体素总数不超过所有训练数据中体素总数的5%。U-Net架构生成的示例在补充说明3的第1和第2节中给出。

8、3D UNet 级联

对下采样数据运行分割模型会增加相对于patch size，从而使网络能够积累更多的上下文信息。这是以减少生成的分割细节为代价的，如果分割目标非常小或其纹理特征，也可能导致错误。在有无限GPU内存的假设场景中，通常倾向于使用覆盖整个图像的patch size来训练全分辨率的模型。3D U-Net级联近似于这种方法，首先在下采样的图像上运行一个3D U-Net，然后训练第二个全分辨率的3D U-Net，以完善前者的分割地图。通过这种方式，“全局”、低分辨率网络使用最大的上下文信息来生成分割输出，然后作为额外的输入通道来引导第二个“局部U-Net”。只有当三维全分辨率U-Net的patch覆盖中值图像形状的12.5%以下的数据集才会触发级联。如果是这种情况，下采样数据的目标间距和相关的3D低分辨率U-Net架构将在迭代过程中共同配置。target spacing被初始化为全分辨率数据的target spacing。为了使patch大小覆盖较大比例的输入图像，在每一步更新结构配置的同时，逐步增加目标间距1%，直到得到的网络拓扑patch大小超过当前图像中值形状的25%。如果当前的间距是各向异性的(最低分辨率和最高分辨率轴之间的差异是2倍)，只有高分辨率轴的间距增加。级联的第二个3D U-Net的配置与独立的3D U-Net相同，其配置过程如前所述(只是将第一个U-Net上采样的分割图连接到其输入)。附注5中的SN5.1b提供了这个优化过程的概述。

经验参数：

Ensembling and selection of U-Net configuration(s).

nnU-Net根据通过对训练数据进行交叉验证计算出的平均前景Dice系数，自动确定使用哪个(集合)配置进行推断。所选型号可以是单个U-Net (2D、3D全分辨率、3D低分辨率或级联的全分辨率U-Net)，也可以是这些配置中的任意两种的集合。模型是通过平均最大概率来集成的。

Post-processing.

基于连通域后处理通常用于医学图像分割。特别是在器官图像分割中，通过去除除最大连通分量外的所有连通分量，往往有助于消除假阳性检测。nnnU-Net遵循这一假设，并自动测试抑制较小连通域对交叉验证结果的影响。首先，所有前景类都被视为一个连通域。如果对除最大区域外的所有区域进行抑制提高了平均前景Dice系数，而没有降低任何类的Dice系数，则选择此步骤作为第一个后处理步骤。最后，nnU-Net基于这一步的结果，并决定是否应该对各个类执行相同的过程。