UNetFormer:一种用于三维医学图像分割的统一视觉 Transformer 模型与预训练框架
图像与视频处理
2022-04-06 v2 人工智能
计算机视觉与模式识别
机器学习
摘要
视觉 Transformer(ViT)近期广受欢迎,因其出色的建模能力,特别是捕捉长程信息的能力,以及对数据集和模型规模的可扩展性,从而在各类计算机视觉与医学图像分析任务中取得最先进性能。本工作中,我们引入一个由两种架构组成的统一框架,称为 UNetFormer,其包含基于 3D Swin Transformer 的编码器以及基于卷积神经网络(CNN)和 Transformer 的解码器。在所提模型中,编码器通过跳跃连接以五种不同的分辨率与解码器相连,并采用深度监督。所提架构的设计能够满足精度与计算成本之间广泛的权衡需求。此外,我们提出一种自监督预训练编码器骨干的方法,通过利用可见 token 的上下文信息来预测随机掩码的体积 token。我们在从公开可用 CT 数据集收集的 张 CT 图像队列上预训练我们的框架,并系统性地研究了掩码比例和块大小等影响表征学习能力与下游任务性能的各种组件。我们使用医学分割十项全能(MSD)数据集对模型进行微调与测试,在肝脏及肝脏肿瘤分割任务上验证了预训练方法的有效性,并在各项分割指标上取得最先进性能。为证明其泛化性,我们使用 MRI 图像在 BraTS 21 数据集上训练并测试模型用于脑肿瘤分割,并在 Dice 分数上优于其他方法。代码:https://github.com/Project-MONAI/research-contributions
引用
@article{arxiv.2204.00631,
title = {UNetFormer: A Unified Vision Transformer Model and Pre-Training Framework for 3D Medical Image Segmentation},
author = {Ali Hatamizadeh and Ziyue Xu and Dong Yang and Wenqi Li and Holger Roth and Daguang Xu},
journal= {arXiv preprint arXiv:2204.00631},
year = {2022}
}
备注
Tech. report, 12 pages, 3 figures