nnUNet 改进 | 让UNet系列在大赛中雄风依旧!

简介: nnUNet 改进 | 让UNet系列在大赛中雄风依旧!

151ec2e25178234a1e9a09b056aebfa7.png

语义分割是医学图像计算中最热门的研究领域之一。尽管 nnU-Net 的概念化可以追溯到 2018 年,但它继续为各种分割问题提供具有竞争力的开箱即用解决方案,并经常被用作挑战获胜算法的开发框架。在这里使用 nnU-Net 参加 AMOS2022 挑战赛,该挑战赛附带一组独特的任务:不仅数据集是有史以来最大的数据集之一,拥有 15 个目标结构,而且比赛还需要提交的解决方案来处理 MRICT 扫描。


通过仔细修改 nnU-Net 的超参数、在编码器中添加残差连接以及自定义后处理策略的设计,能够大大改进 nnU-Net 基线。在提供的训练案例的 5 倍交叉验证中,最终集成在任务 1 (CT) 和任务 2 (CT+MRI) 上的 Dice 得分分别为 90.13 和 89.06。

1、简介

医学图像中所有解剖结构和病理的自动描绘是医学图像计算的长期目标。由于需要专家注释器和 3D 注释的时间密集型特性,到目前为止,数据集需要在目标结构的数量和训练案例的数量之间进行仔细的平衡。因此,现有方法要么在许多图像上进行训练,并且可以稳健地分割少数结构,要么在少数图像上进行训练,可以分割许多鲁棒性降低的结构。因此,无论何时需要对患者解剖结构进行整体透视,都必须将来自多个来源的多个专家模型汇集在一起。这不仅增加了推理时间,而且还产生了新的问题,例如潜在的冲突预测。最后,也许最重要的是,每个专家模型都是独立训练的,标签协同作用不能被利用,因此可能会降低模型的标签效率及其鲁棒性。


在此背景下,2022 年腹部多器官分割 (AMOS2022) 挑战赛旨在促进整体分割方法的发展。它带有数量空前的训练图像和带注释的目标结构:在 500 次 CT 和 100 次 MRI 扫描中标记了 15 个感兴趣的器官,分布在 200+40(CT + MRI)训练、100+20 验证和 200+40 测试图像中挑战提出了两项任务:任务 1 是仅针对 CT 图像的经典多器官分割问题,而任务 2 包括 MRI 图像,并希望提交的方法能够处理这两种模式。


在医学图像分割的背景下,nnU-Net 经受住了时间的考验。尽管具有完全自动化的开箱即用特性,但它始终在新的分割数据集发布时提供最先进的结果。此外,nnU-Net 成功地用作特定任务方法优化的基础,不仅使我们,而且使许多其他团队能够赢得竞争激烈的挑战。因此,使用 nnU-Net 来参与 AMOS2022 挑战似乎是再自然不过的事情了。

2、方法

nnU-Net 是一个自动配置和训练基于 U-Net 的分割管道的框架。通过对目标数据集的严格分析,nnU-Net 自动适应patch sizebatch size、预处理、网络拓扑等。

在本节中,作者对 nnU-net 的自动生成管道提出了一些修改,以最大限度地提高 AMOS2022 挑战的分割性能。在整个方法开发过程中,将所有修改应用于这两个任务,唯一的区别是强度归一化方案。

  • 对于任务 1,使用 nnU-Net 的“CT”方案(数据驱动的裁剪和归一化),
  • 对于任务 2,使用所有图像的简单 z-scorennU-Net 的“非CT”设置)。

2.1、对 nnU-Net 优化

nnU-Net 提供的默认“3d fullres”配置开始,探索了多项改进。尝试用残差编码器替换 U-Net 的默认编码器。进一步优化了预处理,特别是patch sizebatch size和目标间距。GPU 内存容量和处理速度的进步允许比标准 nnU-Net 更大的模型和batch size

67eec80aaf312d8cccf3c02ff85c15dd.png

通过对提供的训练案例运行 5 折交叉验证来执行方法开发。所有模型都是使用 nnU-Net 框架从头开始训练的。实验产生了任务 1 的三个表现良好的候选者和任务 2 的两个候选者,所有这些都总结在表 1 中。

c115d6113e002c6bec7693bde6ad9930.png

图 1 显示了最终配置使用的分割架构。它们共享相同的拓扑,但由于 nnU-Net 自动配置卷积stridekernel-size作为patch size的函数,它们的特征图大小不同。

2.2、推理策略

使用 nnU-Net 默认值(滑动窗口)进行预测。对于验证和测试集预测,使用交叉验证的 5 个模型以及多个不同的配置(任务 1 的 3 个配置:集成中的 3 × 5 = 15 个模型)。集成实现为 softmax 输出的简单平均。

2.3、后处理

nnU-net 提供的后处理旨在涵盖各种用例。作者相信通过分析和定位方法的失败案例,可以在 AMOS2022 上获得额外的性能。为了识别那些,在交叉验证结果上生成了一个混淆矩阵(参见图 2),并对预测的分割图进行了严格的目视检查。

a74c09f32bc99c521ccf5c28e1a1040c.png

  • 左右混乱:有时左肾的一部分被归类为右肾,反之亦然(肾上腺也一样)。在这里,去除不连贯的小组件不是最理想的,因为这样肾脏的一部分就会被标记为背景。因此,将最终的肾脏预测汇集到一个联合类中,并使用连接组件分析来确定连接的肾脏区域。对于每个连接区域,计算其在图像中的位置并相应地分配其标签(左/右)。
  • 连接组件过滤:就像默认的 nnU-Net 一样,探索连接组件的过滤。具体来说,确定是否删除除最大连接组件之外的所有组件可以提高 Dice 分数。
  • 器官大小限制:在人体解剖学的驱动下,器官应该有一定的体积。此属性可用于过滤小的假阳性预测,例如在目标器官不存在但假阳性区域存在的图像中。将一个名为“rate”的参数与在训练集中观察到的最小器官大小相结合。该比率旨在作为安全边际:0.75 的比率表示小于最小器官大小的 75% 的组件被移除。
  • 具有大小限制的组件过滤:将前两种技术结合起来,并删除低于阈值的实例,除非它是唯一的连接区域。
  • 后处理步骤的组成。为了找到最佳的后处理方案,创建了前面提到的后处理方案的组合,即:
  1. PP1:左右混乱(仅肾脏和肾上腺),其次是器官大小限制
  2. PP2:左右混乱(仅肾脏和肾上腺),然后是具有大小限制的组件过滤
  3. PP3:左右混乱(仅肾脏和肾上腺),然后是连接成分过滤。
  4. PP4:只有左右混乱(只有肾脏和肾上腺)。

使用对 5 折交叉验证(训练图像)的集成预测来优化后处理方案。每个器官都进行了独立优化,并保留了性能最佳的后处理策略。请注意,我们会在适用的情况下测试“rate”参数的多个可能值。具体可以参考表 4 和 5 了解最终的后处理配置。

ed8cbe7a63a622672550686aaf4c4bd3.png7598ac3f2f2a3d8b505b9119db5d0b2e.png

3、实验结果

4b62cff971dc22b3ac5272e9fc15e81f.pnge1cc3fd5b4fef4d325b58cf6bbb84462.png

4、参考

[1].Extending nnU-Net is all you need

推荐阅读

华为 + 武汉大学提出基于高斯感受野的标签分配 RFLA,大幅度解决超小目标问题!

BeiT v2 来袭 | BeiT升级,全面超越 MAE,实现 Vision Transformer 微调自由!

Yolov5-6.2 正式发布 | Yolov5 也可以训练分类模型啦,语义分割+实例分割很快到来

相关文章
|
3月前
|
机器学习/深度学习 编译器 TensorFlow
【ASPLOS2024】RECom:通过编译器技术加速推荐模型推理,论文中选并获得荣誉奖项!
2024年5月,关于推荐模型自动编译优化的论文《RECom: A Compiler Approach to Accelerate Recommendation Model Inference with Massive Embedding Columns》在系统领域顶会ASPLOS 2024上中选并进行了展示,并被授予了Distinguished Artifact Award 荣誉,以表彰RECom的易用性与结果的可复现性。
|
6月前
|
存储 编解码 计算机视觉
MobileCLIP来袭 | 如果CLIP可以通过重参加速,你会选择用它作为Backbone预训练吗
MobileCLIP来袭 | 如果CLIP可以通过重参加速,你会选择用它作为Backbone预训练吗
187 0
|
6月前
|
机器学习/深度学习 人工智能 数据可视化
太强!AI没有落下的腾讯出YOLO-World爆款 | 开集目标检测速度提升20倍,效果不减
太强!AI没有落下的腾讯出YOLO-World爆款 | 开集目标检测速度提升20倍,效果不减
656 0
|
机器学习/深度学习 自然语言处理 并行计算
类GPT模型训练提速26.5%,清华朱军等人用INT4算法加速神经网络训练
类GPT模型训练提速26.5%,清华朱军等人用INT4算法加速神经网络训练
163 0
|
开发框架 算法 安全
nnUNet 改进 | 让UNet系列在大赛中雄风依旧!
nnUNet 改进 | 让UNet系列在大赛中雄风依旧!
393 0
|
机器学习/深度学习 编解码 人工智能
谷歌复用30年前经典算法,CV引入强化学习,网友:视觉RLHF要来了?
谷歌复用30年前经典算法,CV引入强化学习,网友:视觉RLHF要来了?
109 0
|
人工智能 算法 网络架构
清华朱军团队开源首个基于Transformer的多模态扩散大模型,文图互生、改写全拿下(1)
清华朱军团队开源首个基于Transformer的多模态扩散大模型,文图互生、改写全拿下
197 0
清华朱军团队开源首个基于Transformer的多模态扩散大模型,文图互生、改写全拿下(2)
清华朱军团队开源首个基于Transformer的多模态扩散大模型,文图互生、改写全拿下
225 0
|
机器学习/深度学习 计算机视觉 网络架构
图神经网络作CV骨干模型,来听听华为诺亚是怎么做的
图神经网络作CV骨干模型,来听听华为诺亚是怎么做的
205 0
|
机器学习/深度学习 人工智能 自然语言处理
150亿参数,谷歌开源了史上最大视觉模型V-MoE的全部代码
150亿参数,谷歌开源了史上最大视觉模型V-MoE的全部代码
570 0