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

本文已在飞桨公众号发布，查看请戳链接：
基于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】公众号
获取更多技术内容~