基于Transformer实现更精准的脑出血多标签分类

原创
06/18 04:38
阅读数 1.7K

本文已在飞桨公众号发布,查看请戳链接:
基于Transformer实现更精准的脑出血多标签分类

图片

灵医智惠是百度旗下深耕医疗领域的AI医疗品牌,多年来一直致力于将AI能力深度赋能医疗行业,加速智慧医疗产业落地。在医疗影像分析领域,灵医智惠也具备了丰富的技术积累,并基于百度飞桨打造了AI影像能力中台EasyMIA(Easy Medical Images Analysis)。目前,EasyMIA已经覆盖眼底影像、CT、MRI、病理切片、皮肤影像等大健康场景下的多种影像分析能力。

背景介绍

颅内出血是严重的健康问题,需要快速的医疗手段介入。脑出血的多标签分类,是脑CT影像分析领域的重要研究方向之一。本次开源的颅内出血多标签分类能力(Intracranial Hemorrhage Detection, IHD)是EasyMIA平台上孵化出的一项业界领先的影像能力。与集成了20个模型的RSNA-IHD竞赛第一名方案相比,百度灵医智惠的模型只需要约1/4的参数和1/10的FLOPs,仅使用单个模型即可实现更精准的脑CT出血分类结果。IHD任务需要为输入的脑CT预测每个切片的出血类别。我们回顾了北美放射学会(RSNA)在 2019 年举办的 IHD 竞赛的top-5解决方案,几乎所有的获胜方案都依赖于2D卷积网络和序列模型(双向GRU或LSTM)来提取切片内和片间特征。在过去几年中,计算机视觉领域取得了很大进展,尤其是基于Transformer的模型在计算机视觉领域大放异彩。因此,我们引入了Transformer来更有效地提取IHD 任务的切片内和切片间的特征。

任务定义

为了鼓励脑出血多标签分类算法的研发,2019年,北美放射学会(Radiological Society of North America, RSNA)在Kaggle竞赛平台上举办了一场国际级算法竞赛(RSNA-IHD)。该竞赛要求参赛选手设计多标签(multi-label)图像分类算法,根据输入的CT切片,判断是否存在颅内出血、及每个具体出血类型。具体来看,根据出血部位的分布,脑出血可以分为5类:分别为硬膜外出血(EDH)、脑实质内出血(IPH)、脑室内出血(IVH)、蛛网膜下腔出血(SAH)、硬膜下出血(SDH)。选手设计的算法需要对每个CT切片,输出一个包含6个元素的概率向量,分别指示切片中是否存在出血和5类具体的出血类型。算法性能通过加权多标签对数损失(log-loss,越低越好)来评估。

图片

图1 脑出血多标签分类及定义

方案选择

榜单中的领先方案

在竞赛中,涌现出许多出色的算法方案。这些方案均首先使用卷积模型提取切片内的图像特征(intra-slice features),接着由序列模型提取切片间特征(inter-slice features)。但是,这些方案存在以下两类问题:

  1. 这些方案在切片内、切片间特征提取阶段,均采用了大量模型集成融合的方案,以提升模型性能。因此,切片内与切片间特征提取过程(模型)是解耦的,无法在GPU上同时训练。这一现象导致了误差累积,限制了解决方案性能的上限。

  2. 大量模型集成融合的方案虽然有利于竞赛成绩的提升,但这类方案因模型数量众多而难以部署至临床环境,推理耗时较长,不利于快速诊断患者的病情。例如,RSNA-IHD竞赛的第一名在切片内特征提取阶段就融合了15个模型,总计参数量超过3.3亿,单次推理运算量超过180G FLOPs。作为对比,我们的方案整体参数量仅为0.86亿,单次推理仅需15G FLOPs。

图片

表1 本文方案与竞赛获胜方案的模型数量与性能指标对比

我们的方案

为解决上述问题,百度灵医智惠影像算法团队基于自研的医疗影像分析平台EasyMIA,研发了基于Transformer的脑CT出血多标签分类方案。该方案在不使用模型集成方案的情况下,单模型即超越了竞赛排行榜第一名(详见表1)。

我们提出一种以端到端方式提取切片内-切片间信息的模型方案。传统的CNN模型擅长提取局部特征,但不能有效地建立全局-长距离图像特征之间的关联。由于脑出血的类型与出血位置密切相关,我们放弃了CNN模型和RNN模型,而是选择改造后的Swin-Transformer网络构建切片内特征提取器、序列Transformer编码器构建切片间特征提取器。受益于Transformer系列模型较高的参数效率,片内和片间的特征提取在我们的解决方案中是完全端到端的,可以在单个现代GPU上运行。这不仅允许梯度信号直接反传到输入层,而且在模型推理期间只需要单次前向传播。

图片

图2 端到端脑出血多标签分类模型架构图

方案详解

模型方案

我们的模型结构如上文图2所示。一组包含N张形状为切片的脑CT序列,选取3种不同的HU窗,经预处理后被组织为的4维张量。切片间特征提取器将N张2D图像变换为N个特征向量,进而在切片间提取器内交换特征信息,最后由分类器(Clas Head 2)输出判别决策。为了避免梯度消失的问题,提升收敛速度,我们对切片内特征连接了一个辅助分类器(Clas Head 1)。在推理时,仅使用Clas Head 2的输出作为模型判别结果。两个分类器均采用Binary Cross-Entropy作为损失函数。

预处理方案

由于颅内出血的类别与出血的位置相关,因此解决方案需要考虑切片中的大脑和颅骨信息。根据竞赛优胜方案的通用设定,我们为每个切片选择3个HU窗(HU值及窗位概念参考[2]),由窗位和窗宽描述:脑组织窗,出血窗,软组织窗。

图片

图3 三种HU窗可视化

训练细节

本方案所有实验均在飞桨2.2.0环境中使用Nvidia-P40 GPU完成。片内与片间特征提取器使用SGD优化器和加权BCE损失函数完成训练,权重为,依序对应“任意”类别和其它5种出血类别。每一轮训练进程持续80,000次迭代(每个CT序列约轮询4次)。在最初300次迭代中,学习率由0线性增加至0.001,并在剩余迭代中沿Cosine曲线逐渐退火至0。

总结

我们提出了一种颅内出血多标签分类方案,仅使用单个模型,其性能即可优于RSNA-IHD竞赛第一名(集成20个模型),单次推理FLOPs减小超过10倍。片内和片间特征提取器以端到端的方式集成到一个模型中,避免了累积误差。但仍存在值得研究的问题,包括更精确的出血多标签分类,以及将所述解决方案扩展到其他切片型应用任务。我们已经发布源代码,希望能够帮助其他研究和工程人员取得更多突破。

更多阅读

  • 技术报告:

https://arxiv.org/abs/2205.07556

  • Github链接:

https://github.com/PaddlePaddle/Research/tree/master/CV/Effective%20Transformer-based%20Solution%20for%20RSNA%20Intracranial%20Hemorrhage%20Detection

  • AI Studio链接:

https://aistudio.baidu.com/aistudio/projectdetail/4047192

参考文献

【1】https://www.kaggle.com/competitions/rsna-intracranial-hemorrhage-detection/overview

【2】https://www.zhihu.com/question/33587197

关注【飞桨PaddlePaddle】公众号
获取更多技术内容~

本文同步分享在 博客“飞桨PaddlePaddle”(CSDN)。
如有侵权,请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”,欢迎正在阅读的你也加入,一起分享。

展开阅读全文
加载中

作者的其它热门文章

打赏
0
1 收藏
分享
打赏
0 评论
1 收藏
0
分享
返回顶部
顶部