STU-Net_stunet

STU-Net: Scalable and Transferable Medical Image Segmentation Models Empowered by Large-Scale Supervised Pre-training

导读

​ 医学图像分割领域的深度学习模型，分为基于 CNN 和 Transformer 两种类型。其中，U-Net 是医学图像分割的先驱 CNN 模型，后续的研究则在此基础上应用残差连接、注意力机制和不同的特征聚合策略等方法。最近，视觉 Transformer 模型被引入到医学图像分割中，例如使用 UNETR 和 SwinUNETR 模型分别采用 Transformer 和 Swin Transformer 作为编码器提取特征。

​ 然而，现有的这些模型不能够适应不同的计算资源和处理不同的医学图像分割任务。此外，虽然大型深度学习模型已经在许多应用领域中表现出非常好的性能，但在医学图像分割领域，最先进的模型仍然非常小，参数数量仅有几千万。

​ 因此，本文提出了一种可扩展和迁移的模型 STU-Net，并探讨了在大规模医学图像分割数据集上训练大型深度学习模型的可能性。

1 引言

​ 医学图像分割是自动标注医学图像中的解剖结构和病变的重要中间步骤，是许多下游临床任务的关键。近年来，各种特定的医学图像分割任务得到了广泛研究，许多基于深度学习的模型取得了巨大的成功。

​ 然而，这些模型通常需要仔细调整以适应不同的任务，这极大地限制了它们的可转移性。因此，需要一种能够同时处理各种医学分割任务的单一模型，包括不同的输入模态（CT、MRI、PET）和不同的分割目标，例如器官和肿瘤。解决这个问题的关键是在大规模数据集上预训练大模型，以使模型具有泛化性。 从数据集的角度来看，一些公共的大规模医学图像分割数据集正在出现。

​ 图1. Total Segmentator 数据集上主流模型的分割性能

​ 此外，大模型通常需要更多的计算成本，特别是当使用三维高分辨率医学图像进行训练时，这一问题会更加严重。因此，本文提出希望这个大模型能够可扩展到不同的尺寸以适应不同的计算预算。

​ 为了实现这一目标，本文提出了一系列可扩展的 U-Net 模型，称为 STU-Net，其参数大小范围从 1400 万到 14 亿不等。此外，为了确保模型的迁移能力，我们使用监督学习的方式在大规模数据集上预训练这些模型。

​ 本文基于 nnU-Net 框架构建了这些模型，因为该框架具有最先进的基线性能，并且被研究人员广泛使用。使用该框架开发大模型存在两个障碍：

1. 基本的卷积块可能不适合扩展。
2. nnU-Net 架构不能被轻易地微调，因为它们被视为超参数，因此具有任务特定性。

​ 为了解决这些问题，本文对 nnU-Net 进行了改进和扩展，提出了一种新的可扩展的大规模医学图像分割模型 STU-Net。此外，在大规模数据集上验证了 STU-Net 的有效性，并在多个下游数据集上展现了出色的泛化性能。

2 相关工作

2.1 医学图像分割模型

​ 医学图像分割以深度学习模型为主，大致分为两类:基于CNN的模型和基于Transformer的模型。U-Net是首创的用于医学图像分割的CNN模型。在此基础上，残差连接、注意力模块和不同的特征聚合策略分别应用于各种任务。

​ 近年来，在自然图像处理中取得成功的具有自注意力机制的Vision Transformer被引入到医学图像分割任务中。其中，UNETR和SwinUNETR分别使用Vision Transformer和Swin Transformer作为编码器，从嵌入的3D patch中提取特征并进行位置嵌入。TransUNet使用Transformer块作为提取全局上下文的瓶颈。nnFormer提出了Transformer和卷积块的交错组合来提取局部和全局特征。这些医学图像分割模型只有几百万个参数，不够大。此外，这些模型不具有可扩展性和可转移性，无法适应不同的计算预算并同时处理不同的医学图像分割任务。

2.2 Scaling-up Models

​ 在深度学习中，按比例放大模型被广泛用于提高模型的性能。最常见的方法是缩放模型的深度和宽度。EfficientNet提出以复合方式扩展网络深度、宽度和分辨率。[20]和[42]分别对语言处理和视觉识别中Transformer的经验缩放规律进行了全面的研究，主要发现计算量、数据大小、模型大小和性能之间的关系符合幂律。根据缩放规律，GPT-3具有1750亿个参数，并在大约45tb文本数据上进行了预训练，在各种文本处理任务中具有接近人类的性能。此外，Vision Transformer已经扩展到220亿个参数，并在大约40亿个图像上进行预训练。使用冻结视觉特征提取器，ViT-22B模型在ImageNet数据集上达到了89.5%的top1精度。然而，将比例模型应用到医学图像的大尺度分割中，却很少有人研究。继EfficientNet之后，[3]对2D U-Net进行生物可降解骨植入物分割。[17]探讨了各种医学图像分割任务的特定于任务的缩放策略。这些工作具有有限的可扩展性，并且仅在小数据集上进行评估。相比之下，我们的论文成功地将模型缩放到比以前的工作大一个数量级，并评估了它们在大规模数据集上的传输能力。

3 方法

​ 我们基于nnU-Net框架构建模型，该框架可以自动配置特定于任务的超参数，并在各种任务上实现最先进的性能。在本节中，我们首先介绍了我们对其进行的改进，以促进可扩展性和可移植性，然后提出了我们提出的缩放方法，最后详细介绍了我们的大规模监督预训练策略，以提高可移植性。

​ 图2。我们的STU-Net体系结构是建立在nnU-Net体系结构的基础上，并进行了一些修改，以增强其可扩展性和可移植性。(a) STU-Net结构概览。蓝色箭头表示下采样，黄色箭头表示上采样。(b)残差块，实现大尺度模型。©每个编码器级的第一个残差块的下采样。(d-e)输入输出信道转换的干和分段头。(f)上采样的无权重插值，有效解决不同任务间的权重不匹配问题。

3.1 nnU-net架构

​ nnU-Net 采用了基于跳跃连接的对称编码器-解码器体系结构，其中包含各种分辨率阶段。每个阶段由两个卷积层组成，后跟 Instance Normalization 和 Leaky ReLU (Conv-IN-LeakyReLU)。

​ 由于其不包含残差连接，简单地在每个阶段堆叠更多层可能会遭受梯度扩散，使整个模型难以优化。这限制了 nnU-Net 的深度，并进一步限制了其可扩展性。

​ 另一方面，nnU-Net 根据数据集属性确定输入 patch 大小和输入间距。然后，使用数据集特定的 patch 大小和间距来设置与网络体系结构相关的超参数，例如分辨率阶段的数量、卷积核和下采样/上采样比率。因此，这些与体系结构相关的超参数在任务之间有所不同，导致不同的任务具有不同的网络体系结构。此外，训练在一个任务上的模型不能直接迁移到其他任务中，这限制了模型的传递能力评估。

3.2 基于 nnU-Net 改进

​ nnU-Net 的任务特定的超参数可以分为与模型权重相关的 (如卷积核大小、分辨率级数) 和与模型权重无关的 (例如池化核大小、输入图像块大小和间距等)。

​ 为了使模型架构更适合迁移到其他任务中，我们固定了与模型权重相关的超参数，即将所有任务的分辨率级数保持为 6，并对所有卷积层使用等向性的卷积核(3，3，3)。对于与模型权重无关的超参数，我们采用了 nnU-Net 的默认设置，以确保其在各项任务上具有最先进的性能。本文还将我们的设置与 nnU-Net 和 3D U-Net 进行了比较。

3.2.1 基本块的改进

​ nnU-Net 的每个阶段都由一个基础块构成，每个基础块由 两个Conv-Instance Normalization- LeakyReLU 层组成。但当增加每个阶段中的基础块数量时，会由于梯度扩散而出现优化问题。

​ 为了解决这个问题，我们在基础块中引入了残差连接。除了第一个分辨率阶段，编码器中的每个阶段都从图2©中的下采样块开始，然后是图2(b)中的几个残差连接块。此外，为了使整个架构更加紧凑，我们还将降采样集成到每个阶段的第一个残差块中。这种降采样块具有与常规残差块相似的残差架构，由左、右两个分支组成，如图2©所示其中左分支有两个不同步长的 3×3×3 卷积，而右分支使用步长为 2 的 1×1×1 卷积核。该下采样块与图2(b)中的规则残差块具有相似的残差架构，这种基础块的改进使得整个架构更加紧凑，同时还能解决梯度扩散的问题。

3.2.2 上采样调整

​ nnU-Net 的上采样默认使用转置卷积 (transpose convolution) 进行。但对于不同任务，卷积核和步长可能会在相同的分辨率阶段内变化，这会导致转置卷积的权重形状不同，从而使得在不同任务之间进行权重传递时出现权重不匹配的问题。

​ 如图2(f)，为了解决这个问题，我们使用插值 (interpolation) 加上一个步长为 1 的 1×1×1 卷积层来替代转置卷积。这种权重自由的插值方法可以解决权重形状不匹配的问题。我们使用最近邻插值 (nearest neighbor interpolation) 来进行上采样，实验结果（表6）表明，最近邻插值不仅速度更快，而且还能够达到与双立方插值 (cubic linear interpolation) 相当的性能。

3.2.3 缩放策略

​ 深度网络通常具有更大的感受野和更好的表示能力，而宽网络则往往在每个层中提取更丰富的多尺度特征。根据 EfficientNet 的研究结果显示，深度缩放和宽度缩放不是独立的，为了达到更好的准确性和效率，最好以复合方式缩放网络的深度和宽度。

​ 为了简化缩放问题，我们采用了对称结构的模型，即同时缩放编码器和解码器，并在每个分辨率阶段中以相同的比例缩放深度和宽度。表2展示了 STU-Net 的不同规模，其中后缀“S，B，L，H”分别表示 “Small, Base, Large, Huge”。

3.2.4 预训练

​ 我们使用 Total Segmentator 数据集对 STU-Net 进行了预训练，STU-Net 最终的 1×1×1 卷积层有 105 个通道，对应于 Total Segmentator 中目标注释类别的总数。

​ 为了使预训练模型更加通用和可转移，我们对 nnU-Net 中的标准训练过程进行了一些修改。与 nnU-Net 中默认的 1000 个训练 epoch 相比，我们将模型预训练了 4000 个 epoch。此外，我们发现使用镜像数据增强可以提高模型在下游任务上的转移性能。

​ 预训练模型可以直接对由 CT 图像组成且包含上游 104 个类别目标分割类别的下游数据集进行直接推理，无需作进一步的调整。

​ 对于具有新标签或不同模态的下游任务，我们使用训练好的模型作为初始化，并随机初始化分割输出层以匹配目标输出类别的数量。在微调过程中，分割头部是随机初始化的，而其余层的权重则从预训练模型中加载。这些权重使用比分割头部更小的学习率 (0.1倍) 进行微调，从而获得更好的结果。

4 实验

数据集：

​ 我们在TotalSegmentator[38]数据集上训练不同尺度的STU - Net，该数据集包含1204张图像，包含104个解剖结构(包括27个器官，59个骨骼，10个肌肉和8个血管)。它涵盖了大部分临床对全身正常结构的分割目标。所有图像重新采样到1.5 × 1.5 × 1.5 mm各向同性分辨率。我们遵循[38]中的原始数据分割，其中使用1081个案例用于训练，57个案例用于验证，65个案例用于最终测试。值得注意的是，由于数据隐私原因，人脸被模糊了。我们在14个公共数据集上评估了我们训练好的STU-Net，用于直接推理和3个公共数据集，用于进一步微调，以测试我们训练模型的可移植性。这些下游数据集的详细属性见附录。

评价指标：

​ 我们采用Dice Similarity Coefficient (DSC)作为评价指标，DSC分数越高表示分割效果越好。为了公平的比较，我们报告了模型在最后一个epoch训练的结果，而不是最好的一个epoch。

实现细节：

​ 我们在Python 3.8、CentOS 7、Pytorch 1.10和nnU-Net 1.7.0环境下运行所有实验。我们大致遵循nnU-Net中默认的数据预处理、数据增强和训练过程。我们使用Nestrov动量为0.99的SGD优化器，权重衰减为1e-3。批大小固定为2，每个epoch包含250次迭代。

​ 对于所有数据集，从头开始训练时，学习率从0.01开始，除了AutoPET，它从0.001开始，遵循最先进的解决方案[40]。学习速率按照poly学习率策略(1−epoch/1000)0.9衰减。在训练过程中，我们采用了加性亮度、伽马、旋转、缩放、镜像和弹性变形等动态数据增强方法。TotalSegmentator上的预训练补丁大小为128×128×128。

​ 下游任务的微调补丁大小由nnU-Net自动配置。模型在具有80 GB VRAM的NVIDIA Tesla A100卡上进行训练。

4.1 TotalSegmentator的定量结果

​ STU-Net-B 模型在所有类别的平均 DSC 方面分别超过基于 CNN 的最佳模型 nnU-Net 和基于 Transformer 的最佳模型 SwinUNETR-B 0.36% 和 4.48%。

​ 将我们的基础模型进一步扩展到大尺寸和超大尺寸分别导致平均 DSC 分数提高 1.59% 和 2.94%。

​ STU-Net-H 在 Total Segmentator 数据集中的所有类别和五个子类组中实现了最高的平均 DSC。 结果显示了我们对 nnU-Net 和扩展策略的架构改进的有效性。

4.2 训练模型的可转移性

​ 我们通过以下方式评估训练模型的可转移性:1)对下游CT数据集进行直接推断，这些数据集包TotalSegmentator数据集中104个类别的子集;2)对三个下游数据集上的训练模型进行微调，包括TotalSegmentator数据集中不存在的类别(例如病变)和CT以外的模式(例如MR, PET)。

4.2.1 直接推理结果

​ 使用 Total Segmentator 进行预训练时，规模较大的模型通常在所有这 14 个数据集中具有更高的平均 DSC 分数。

4.2.2 微调结果

与在下游数据集上从头开始训练的模型相比，微调我们在 Total Segmentator 上预训练的 STU-Net 模型可以获得更好的分割性能。

可以很直观地看出，STU-Net 的分割结果在完整度和细腻度上都比其他的模型表现更好，这更充分地证明了 STU-Net 在医学图像分割领域中的先进性与通用性。

4.3 消融实验

​ 我们提出的STU-Net-B模型是默认的nnU-Net体系结构的改进版本，包含了几个改进。表6比较了不同STU-Net-B体系结构变体在TotalSegmentator验证集上的分割性能。我们的STU-Net-B模型利用最接近的插值进行上采样，并在第一个残差块中合并下采样。为了评估其性能，我们对使用单独卷积的替代下采样方法和使用转置卷积或三线性插值的替代上采样方法进行了比较分析。

​ 我们首先将标准nnU-Net的最大特征数从320增加到512，以匹配我们的STU-Net-B的参数，并将其记为nnU- Net*。nnUNet *比标准的nnU-Net工作得更好，但比我们的STU-Net-B性能略差。这个比较证明了在STUNet-B中改进的有效性。

然后，我们通过引入一种采用卷积下采样的变体来探索STU-Net中的下采样设计，而不是在每个阶段的第一个残差块内集成下采样过程。这种修改会导致性能下降和计算成本增加。通过设计两种不同版本的STU-Net-B，我们进一步研究了上采样的改进，第一种是利用转置卷积来代替默认的插值和基于卷积的上采样。这将导致性能下降0.16%，并且使权重不能用于下游微调。第二种方法使用三线性(或三次线性)插值来代替最近邻插值。这种变化降低了性能，减慢了运行速度。总的来说，STU - net默认的上采样设计实现了更好的性能、更快的运行速度和更好的传输容量。

因此，我们提出的改进不仅提高了nnU-Net的有效性和效率，而且赋予了它权值可转移性和可扩展性的关键特性。这些属性对于进一步扩展模型和促进迁移学习至关重要

4.3.2 缩放策略

我们对nnU-Net和我们的STU-Net-base采用了三种不同的缩放策略，即使用不同的深度系数d∈[1.0,2.0,3.0,4.0]，宽度系数w∈[1.0,2.0,3.0,4.0]，同时使用深度和宽度系数[1.0,2.0,3.0]进行缩放。系数d(或w)表示每个阶段的深度(宽度)缩放为d(或w)倍。表7显示了TotalSegmentator数据集上不同缩放策略的结果。首先，较宽的nnUNet*和STU-Net始终获得更好的性能，而较深的则没有。因此，与深度缩放相比，宽度缩放更能有效地提高模型在大规模数据集上的性能。但它带来了计算消耗的显著增加。其次，与nnU-Net相比，性能下降更深中，我们的STU-Net在一定程度上可以更好地受益于缩放深度，例如87.12→87.72→87.99→87.58 vs. nnUNet的86.94→85.65→83.70→81.45。STU - Net的优良性能得益于它的残差设计。第三，复合缩放，即同时增加深度和宽度，比其他两种缩放策略更有效和高效地提高了我们的STU-Net的性能。最后，即使使用相同的扩展策略，以及相似的参数和FLOPs，我们的STU-Net在所有设置上始终优于nnU-Net *，这再次验证了我们改进的有效性。

4.3.3 预训练和微调策略

​ 我们对预训练的有效性进行了实证研究，然后对在预训练和微调阶段使用的镜像增强进行了实证研究。最后，研究了预训练的训练时段。

​ 首先，我们研究了大规模预训练的有效性。我们使用STU-Net-L在TotalSegmentator数据集上进行预训练，然后在FLARE22和AutoPET数据集上进行微调来进行这样的研究。对比表8中第1行(w/o预训练)和第3行(w/o预训练)，我们发现经过预训练的模型比未经过预训练的模型分别高出0.84%和3.67%。这种更好的性能证明了大规模预训练的有效性。

​ 其次，我们研究了镜像增强在预训练和微调中的应用。通过预训练，镜像有助于提高下游任务的DSC，因为镜像(第3 ~ 5行)大大提高了第6行(无镜像)的DSC。同样值得注意的是，在预训练和微调阶段使用的镜像都达到了最好的效果。

​ 最后，将镜像结合到预训练和微调中，我们发现不同的预训练时期也会影响性能(参见最后四行)。在4k epoch得到了最好的结果，可以保证模型得到充分的收敛训练。

4.3.4 数据集大小对模型性能的影响

​ 在TotalSegmentator数据集上训练模型时，我们研究了数据集大小对模型性能的影响。值得注意的是，不同比例的训练案例是通过分层随机选择过程获得的，以确保较高比例的训练案例也包括较低比例的数据。

​ 如图4所示，无论训练用例的数量如何，增加模型大小都会导致TotalSegmentator子集上更好的分割性能。例如，即使只训练了5%的案例，STU-Net-H的表现也优于STU-Net-S。同样地，当只训练20%的案例时，STU-Net-H超过了STU-Net-B。这些结果表明，在医学图像分割中，大规模模型比小型模型具有更高的数据效率。此外，随着案例数量的增加，不同型号的性能也在不断提高，趋势尚未饱和。这些观察结果表明，基于TotalSegmentator数据集增加训练案例的数量可以进一步提高性能。

4.4. 通用模型与专家模型

​ 我们评估了在TotalSegmentator数据集中所有104个类上训练的通用STU-Net模型的性能，对比了五个专家su - net模型，每个模型针对五个子类别中的一个(与表3相同)。此外，我们研究了模型大小对专家模型和通用模型性能的影响。

​ 图5表明，随着模型大小的增加，专家模型和通用模型的性能通常都会提高。专家模型在器官、椎骨和心脏亚类别中表现出色，而通用模型在肌肉和肋骨亚类别中表现更好。对于最大的模型(STU-Net-H)，通用模型超过了专家模型，在TotalSeg数据集的所有类别上实现了最高的总体平均DSC得分为90.06%，而专家模型的最高平均DSC得分为89.07%。

​ 结果表明，尽管专家模型在特定的子类别中可能优于通用模型，但通用模型在不同的解剖结构中始终具有较强的性能。随着模型大小的增加，专家模型和通用模型之间的性能差距在不同的子类别中有所不同:器官和椎骨的差距缩小，心脏的差距逆转，肌肉和肋骨的差距扩大。研究结果表明，随着模型尺寸的增加，通用模型能够同时分割多个类别，并表现出有希望的性能进步。

5 结论

​ 本文介绍了一种基于 nnU-Net 框架的可扩展的医学图像分割模型 STU-Net。STU-Net 最大包含 14 亿个参数，是迄今为止最大的医学图像分割模型。通过在大规模的 Total Segmentator 数据集上训练 STU-Net 模型，我们证明了模型规模的扩展在迁移到各种下游任务时产生了显着的性能提升，同时这验证了大模型在医学图像分割领域的潜力。

​ 此外，STU-Net-H 模型在 Total Segmentator 数据集上训练，在多个下游数据集中表现出强大的直接推理和微调能力。这一观察结果强调了利用大规模预训练模型进行医学图像分割任务的实际价值。

​ STU-Net 模型的发展有望推动医学图像分割技术的发展，为医学图像分割社区的研究和创新开辟了新的途径。