STU-Net: 推动医学图像领域走向 ImageNet + ResNet 时代

医学图像领域的发展相较于自然图像领域存在着明显的滞后。自然图像发展的一个里程碑是建立了经过 ImageNet 数据集预训练的不同大小的 ResNet 分类模型，随后迁移到更多的任务上。遗憾的是，医学图像领域至今还没有完全走到 ImageNet + ResNet 时代，其根本原因是缺乏大规模的有标注数据集以及标准化的可迁移的不同大小模型。医学图像分割是医学图像分析中最基础的任务，而 U-Net 是医学图像分割中最基础有效的模型。研究者认为，在大规模医学图像分割数据集上预训练不同大小的标准化 U-Net 可以推动医学图像分析领域的快速发展。
该文章提出了一系列可扩展且具有良好迁移能力的U-Net模型——STU-Net，其参数规模范围从1400万到14亿。值得注意的是，具有14亿参数的STU-Net是目前最大的医学图像分割模型。他们的 STU-Net 基于 nnU-Net 框架，由于其应用广泛和出色的性能。他们首先优化了 nnU-Net 默认的卷积块，使它们具有可扩展性。然后，我们经验性地评估了不同的网络深度和宽度的扩展组合，发现同时扩大模型的深度和宽度是最优的。首先，他们在大规模的 TotalSegmentator 数据集上训练不同大小的STU-Net模型，并发现增加模型大小可以带来更强的性能。这意味着在医学图像分割领域，大规模模型有着巨大的潜力。此外，他们评估了模型在14个下游数据集的直接推理及3个数据集的微调性能，包括不同的图像模态和分割任务。结果显示，文章提出的预训练模型无论是在直接推理还是微调上都有出色的表现。

论文:

https://arxiv.org/abs/2304.06716

开源代码（点击“阅读原文”即可直达）:

https://github.com/uni-medical/STU-Net

nnU-Net不具有迁移性和扩展性

nnU-Net 是目前性能最先进的医学图像分割框架，直到最新的 MICCAI 2023 各类医学图像分割竞赛，冠军方案仍大多是基于 nnU-Net 框架。nnU-Net 很好的考虑到了不同医学图像任务和数据集之间存在的巨大差异，自适应的去配置数据预处理策略和模型结构来训练出性能更好的模型。这些数据集特异性的设计在给 nnU-Net 带来了先进性能的同时，却也导致了 nnU-Net 在一个数据集上训练的结构特异性的模型，很难直接迁移到其它数据集上。另外一方面，nnU-Net 的模型大小相对固定，其结构中的基本卷积块不适合扩展，也很少有工作指导，如何去合理的扩展医学分割模型的大小。

STU-Net改进nnU-Net的

迁移性和扩展性

图1：STU-Net 结构示意图，基于 nnU-Net 对一些模块进行修改，使其具有迁移性和可扩展性。

超参数的设置

nnU-Net 的任务特定的超参数可以分为与模型权重相关的 (如卷积核大小、分辨率级数) 和与模型权重无关的 (例如池化核大小、输入图像块大小和间距等)。为了使模型架构更适合转移到其他任务中，我们固定了与模型权重相关的超参数，即将所有任务的分辨率级数保持为 6，并对所有卷积层使用各向同性的卷积核 3×3×3。对于与模型权重无关的超参数，我们采用了 nnU-Net 的默认设置，以确保其在各项任务上具有最先进的性能。相比完全固定的 3D U-Net 和过于任务特异性的 nnU-Net，我们的超参数设置在可迁移性和模型性能之间取得了很好的平衡。

图2：对比 STU-Net 和 nnU-Net, 3D U-Net 之间的超参数设置。

调整基本模块

nnU-Net 的每个阶段都由一个基础块构成，每个基础块由两个Conv-Instance Normalization- LeakyReLU 层组成。但当增加每个阶段中的基础块数量时，会由于梯度扩散而出现优化问题。
为了解决这个问题，我们在基础块中引入了残差连接。此外，为了使整个架构更加紧凑，我们还将降采样集成到每个阶段的第一个残差块中。这种降采样块具有与常规残差块相似的残差架构，由左、右两个分支组成，其中左分支有两个不同步长的 3×3×3 卷积，而右分支使用步长为 2 的 1×1×1 卷积核。这种基础块的改进使得整个架构更加紧凑，同时还能解决梯度扩散的问题。

调整上采样模块

nnU-Net 的上采样默认使用转置卷积 (transpose convolution) 进行。但对于不同任务，卷积核和步长可能会在相同的分辨率阶段内变化，这会导致转置卷积的权重形状不同，从而使得在不同任务之间进行权重传递时出现权重不匹配的问题。
为了解决这个问题，我们使用插值 (interpolation) 加上一个步长为 1 的 1×1×1 卷积层来替代转置卷积。这种权重自由的插值方法可以解决权重形状不匹配的问题。我们使用最近邻插值 (nearest neighbor interpolation) 来进行上采样，实验结果表明，最近邻插值不仅速度更快，而且还能够达到与双立方插值 (cubic linear interpolation) 相当的性能。

缩放策略

深度网络通常具有更大的感受野和更好的表示能力，而宽网络则往往在每个层中提取更丰富的多尺度特征。根据 EfficientNet 的研究结果显示，深度缩放和宽度缩放不是独立的，为了达到更好的准确性和效率，最好以复合方式缩放网络的深度和宽度。
为了简化缩放问题，我们采用了对称结构的模型，即同时缩放编码器和解码器，并在每个分辨率阶段中以相同的比例缩放深度和宽度。图 3 展示了 STU-Net 的不同规模，其中后缀“S，B，L，H”分别表示 “Small, Base, Large, Huge”。

图3：不同大小的 STU-Net 的深度宽度设置和对应的参数量和计算量。

基于 TotalSegmentator

大规模有监督预训练 STU-Net

我们选择了 TotalSegmentator 数据集对 STU-Net 进行了预训练，他包含了1204张三维图像和104种类别，其中覆盖了27类器官、59类骨骼、10类肌肉和8类血管，涉及了全身各个部位的常见临床分割目标。为了使预训练模型更加通用和可迁移，我们对 nnU-Net 中的标准训练过程进行了一些修改。与 nnU-Net 中默认的 1000 个训练 epoch 相比，我们将模型预训练了 4000 个 epoch。此外，我们发现使用镜像数据增强可以提高模型在下游任务上的转移性能。预训练模型可以直接对由 CT 图像组成且包含上游 104 个类别目标分割类别的下游数据集进行直接推理，无需作进一步的调整。对于具有新标签或不同模态的下游任务，我们使用训练好的模型作为初始化，并随机初始化分割输出层以匹配目标输出类别的数量。在微调过程中，分割头部是随机初始化的，而其余层的权重则从预训练模型中加载。我们发现，相较于 nnU-Net 从头训练的学习率，使用更小的学习率微调，具有更好的效果。

实验结果

在大规模数据集上，更大的模型具有更好的性能

我们首先在 TotalSegmentator 官方训练和验证数据集上的上训练和测试模型，对比 STU-Net 和其它常用的医学图像分割模型。可以观察到，更大的模型具有更好的性能。CNN 模型要优于 Transformer 的模型。具有 14 亿参数量的 STU-Net-H 模型取得了最好的性能。

图4：在 TotalSegmentator 上对比不同模型的性能。气泡的大小与模型的计算量相关，不同的颜色代表不同的模型，同一种颜色的不同大小的气泡对应同一种模型的不同尺寸。

图5：不同模型在 TotalSegmentator 数据集上五个子类别和所有类别上平均性能对比。

经过大模型预训练，更大的模型具有更好的泛化性

我们将在 TotalSegmentator 上预训练的模型直接在 14 个不同的下游数据集上直接推理。实验结果表明，更大的模型具有更好的泛化能力。

图6：大规模数据集上，更大的模型具有更好的跨数据集泛化性。

大规模有监督预训练有效，更大的模型具有更好的迁移性

我们在 3 个数据集的 7 个子任务上，对经过 TotalSegmentator 预训练的不同的大小的 STU-Net 进行微调。相比从头训练的 STU-Net 以及 nnU-Net，经过预训练的 STU-Net 表现更好，而且模型越大，其迁移性越好。

图7：经过 TotalSegmentator 预训练的模型，在 7 个子任务上具有良好的迁移性。

图8：可视化不同模型的分割结果。经过预训练的更大规模模型，具有更好的分割效果。

更大的模型具有更高效的数据利用率

我们对比了使用不同比例的 TotalSegmentator 训练数据训练的不同大小的 STU-Net，更大的模型具有更好的数据利用率，使用少量的训练数据即可达到小模型使用大量数据训练的性能。

图9：对比不同大小的模型在不同比例训练数据下的性能。更大的模型只需要少量训练数据，即可取得与小模型在大量训练数据下的性能。

更大的模型更适合用作通用模型

总结

本文介绍了一种基于 nnU-Net 框架的可扩展和可转移的医学图像分割模型 STU-Net。STU-Net 最大包含 14 亿个参数，是迄今为止最大的医学图像分割模型。通过在大规模的 TotalSegmentator 数据集上训练 STU-Net 模型，我们证明了模型规模的扩展在迁移到各种下游任务时产生了显著的性能提升，同时这验证了大模型在医学图像分割领域的潜力。此外，STU-Net-H 模型在 Total Segmentator 数据集上训练，在多个下游数据集中表现出强大的直接推断和微调可迁移性。这一观察结果强调了利用大规模预训练模型进行医学图像分割任务的实际价值。总之，可扩展和可转移的 STU-Net 模型的发展有望推动医学图像分割技术的发展，为医学图像分割社区的研究和创新开辟了新的途径。

END

编辑 | 刘帅

来源 | 通用医疗GMAI

审核 | 医工学人

医工学人简介

医工学人是在医疗科技创新与医工交叉背景下成立的多高校学生学术组织。旨在建立医学、工程学领域研究者的对话渠道，创造交流分享医工交叉前沿技术的优质平台，推动医疗科技创新与医工交叉融合。

目前组织内共有六百多位来自复旦大学、西安交通大学、上海交通大学、清华大学、浙江大学、中国科学技术大学、各高校附属医院等30余所重点高校、医院及科研单位医工学科相关的医生、学生、专家学者等。欢迎志同道合的你加入我们！

<扫码添加社群管理人微信>

本篇文章来源于微信公众号: 医工学人

基于 TotalSegmentator

大规模有监督预训练 STU-Net

发表回复 取消回复

发表回复取消回复