基于 CPU 实施高效且快速的医疗影像分析

基于 2017 年阿里巴巴天池医疗 AI 大赛的实验性总结

作者:胡潇吴慧姚伟峰

概述

本文基于由阿里云、英特尔和零氪科技联合主办的天池医疗 AI 大赛。在大赛中选手挑战通过计算机断层 (CT) 扫描自动检测肺结节。该赛事吸引了全球 2,887 支参赛队伍,经过 7 个月的紧张角逐,在 2017 年 10 月成功落下帷幕。大赛托管于阿里巴巴的公有云服务,此服务完全基于英特尔的深度学习软硬件堆栈而构建。在整个比赛过程中,英特尔积极参与了架构设计、软硬件开发、性能优化和在线支持,因此积累了有关医疗 AI 的丰富洞察。本文将以实验方式阐述这些重要发现。

首先,我们吸收了所有天池参赛者共同的设计原则,实现了一个具有代表性的 3D 卷积神经网络 (CNN) 模型,以反映当前最新的肺结节检测水平。

第二,我们运用分辨率不同的输入数据训练该模型,并以定量的方式证明,用分辨率较高的数据训练的模型可有效提升检测性能,尤其是针对小结节的检测性能。但同时,高分辨率模型所使用的内存比低分辨率模型多。

第三,我们对通用图形处理器 (GPGPU) 和 CPU 的行为进行了比较,并证明 CPU 架构可提供更大的内存容量,因此有助于医疗 AI 开发人员探索高分辨率设计,从而实现最佳检测性能。

我们还引进了定制的深度学习框架,名为“ Extended-Caffe*”[10],它是天池软件堆栈的核心,我们以此为例有效地证明了 CPU 架构可支持快速 3D CNN 计算,进而方便人们高效且快速地在CPU上进行3D CNN 模型开发。

背景

天池医疗 AI 大赛 [1] 是中国首次举办的 AI 医疗大赛,其规模和数据量在世界范围内均属首例。中国 16 家顶尖肿瘤医院为大赛提供近 3,000 例患者的肺部 CT 扫描数据。之所以选择肺结节检测为主题,是因为过去 30 年以来,中国肺癌患病率逐年上升,已经成为癌症中死亡率最高的疾病。因此,解决肺结节早期筛查问题尤为重要,刻不容缓。全球 2,887 支参赛队伍经过为期 7 个月的紧张比赛,最终北京大学团队脱颖而出,获得冠军。

在线比赛托管于阿里巴巴的公有云服务,此服务完全基于英特尔的深度学习软硬件堆栈而构建。基础硬件设施为基于英特尔® 至强® 和英特尔® 至强融核™ 处理器的平台集群,该平台提供总共 400+ TFLOPS 的计算能力。英特尔还提供一系列深度学习软件组件支持模型训练和推理,其核心是定制的深度学习框架,称为 Extended-Caffe,专为医疗 AI 用途而优化。英特尔专家还帮助数百位在线参赛者高效运行模型,在此过程中积累了有关医疗 AI 领域的重要洞察。

此次大赛表明,尽管深度学习应用于计算机视觉领域已超过十年,但对于领域专家和工程师来说,医疗影像分析仍然面临着严峻的挑战。特别是几乎所有最先进的医学影像分析解决方案都严重依赖于 3D甚至 4D/5D CNN,相比于计算机视觉领域的其他常见 2D CNN,它们在工程设计方面的考虑因素有很大的不同。我们发现,相比于传统 GPGPU,CPU 平台可更有效地支持使用 3D CNN 进行医疗影像分析,因为 CPU 在内存容量方面具有显著的优势。而同时,通过巧妙的底层算法设计,CPU平台可以支持快速的3D CNN元语操作,以保证模型训练和推理的速度。

尽管天池数据集和模型不对外公开,但我们自主研发和试验了一个 3D CNN 模型,以反映当前最新的肺结节检测结果,以此阐述我们的重要结论。下文将逐一介绍我们如何预处理 CT 数据集、如何设计模型、在 CPU 上如何实现高速 3D元语操作、以及实验步骤,并通过定量分析得出结论。

CT 数据预处理

原始 CT 影像包含一系列 2D 切片。相邻2D 切片之间的间隔称为Z间隔。每张2D切片都是一个灰度像素矩阵,其中像素之间的水平间隔和垂直间隔分别称为 X 间隔和 Y 间隔。原始的CT影像数据均记录这些间隔的实际长度(以毫米为单位)。由于 CT 仪器之间经常存在差异,因此不同的原始 CT 影像其间隔也各不相同。例如,在 LUNA'16 数据集 [3] 中,Z 间隔从 0.625 毫米到 2.5 毫米不等。天池数据集也存在类似的情况。为了使深度学习模型所处理的数据集保持统一,我们必须以固定的采样距离,在 X、Y 和 Z 三个方向对原始像素进行插值,从而使原始 CT数据 转换为新的3D图像,其中3 个方向的像素间隔均等于该采样距离。然后,以像素为单位进行测量,新的3D图像和结节的大小(即分辨率)将取决于该采样距离。表 1 显示较小的采样距离可提高分辨率。注意,与检测普通物体不同,检测肺结节时遇到的特有问题是,结节仅占整个 CT体积的百万分之一。因此,为了使模型能有效提取结节的特征,我们必须将 CT 图像裁剪成较小的 3D 区域,然后将这些裁剪区逐一导入模型中进行训练。同样,较小的采样距离对应较大的裁剪区。

表 1.不同的采样距离导致分辨率不同的 3D 数据

采样距离 (毫米)3D 图像分辨率 (像素 x 像素 x 像素)结节分辨率 (直径:像素)裁剪区分辨率 (像素 x 像素 x 像素)
1.00249x256x3023.66128x128x128
1.33188x196x2312.7496x96x96
2.00127x136x1591.8364x64x64

用于肺结节检测的 3D CNN 模型

图 1.我们的 3D CNN 模型架构(以输入的裁剪区大小 = 128 x 128 x 128为例)

根据业界其他成果 [4‒6] 和天池大赛中的各种模型,参考他们的共同原则,我们构建了一个用于肺结节检测的 3D CNN 模型,如图 1 所示,该模型分为下采样和上采样两部分。下采样部分由 5 个 3D 残差块与 4 个池化层相互交织而成。残差块由卷积、批量归一化、ReLU 和其他操作组成,另外还有一个残差结构(图 1 中的 C1 和 C2)。上采样过程通过两个反卷积(图1中的 Deconv)来完成。我们将每个反卷积的输出和相应下采样层的输出进行合并,使所得的特征图包含原始输入数据的局部和全局信息。

对于每个输入裁剪区 (m x m x m),我们的模型生成 (m/4) x (m/4) x (m/4) x 3 个立方体区域框,称为“候选区”。为每个候选区会生成一个概率值(即该候选区内包含结节的可能性),一组坐标(用以表征该候选区的中心点),以及该候选区的大小(即边长)。

通常,在3D CNN模型之后,仍需要进行一系列的后处理,包括假阳性剔除及一些其他步骤,用以进一步筛选出真阳性候选区。不过,由于3D CNN模型本身的有效性和速度是工程实施的决定性因素,因此,本文没有关注这些步骤。但是,即便没有辅以这些后处理步骤,我们所训练的模型(CCELargeCubeCnn [9] )依然在LUNA’16排行榜中排名第 14位,这足以说明我们的这一实验模型的确代表了当前最先进的肺结节检测水平。

CPU 上的高效 3D CNN 元语操作

我们在实验中通过使用不同的分辨率和超参数(例如批处理规模)运行模型,对比 GPGPU 和 CPU 平台的有效性。因此,必须首先确保在 CPU 平台上的运算速度,特别是针对最常用的3D 卷积元语操作。

CPU 上的 3D 卷积元语操作

我们利用面向深度神经网络的英特尔® 数学核心函数库(英特尔® MKL-DNN [7])中高度优化的 2D 卷积元语,在 CPU 上实现了高效的 3D 卷积元语(见图 2)。首先,我们将 3D 数据和卷积核分别视为一组 2D 切片。然后,一个3D 卷积运算就等价于先对相应的 2D 切片(图 2 中颜色相同的切片)分别进行卷积操作,然后再将所有的中间结果相加。由于英特尔® MKL-DNN的 2D 卷积元语在 CPU 上已经经过了高度的优化,因此我们的 3D 卷积也可以在 CPU 上高速运行。

图 2.快速 3D 卷积元语的实现原理(调用MKL-DNN的2D卷积元语)

为了显示效果,我们还开发了另一个基准实现,称为基于普通矩阵乘法 (GEMM) 的 3D 卷积。它是一个简单直接的实现。图 3 描述了一个基于GEMM 的2D卷积算法的工作原理,即重新排列数据和卷积核,然后使用矩阵乘法计算出卷积结果。而在这里,我们运用相同的方法处理 3D 数据和卷积核,就得到了基于GEMM 的3D卷积。因为这种算法的核心运算是矩阵乘法,对此我们可以利用英特尔® 数学核心函数库 (英特尔® MKL) [8] 中高度优化的 SGEMM 来实现。也就是说,我们的这一基准实现实际上已经可以在 CPU 上达到较高的速度。

图 3.基于 GEMM 的 2D 卷积(利用英特尔® 数学核心函数库 SGEMM)

而图 4 则显示了我们实现的快速 3D 卷积元语的速度仍大大高于基于 GEMM 的3D卷积元语基准实现:前向路径的速度提升了 4 倍,后向提升了 30%,而整体速度提升了 2 倍。

图 4.执行时间比较(快速 3D 卷积元语实现对比于基于 GEMM 的 3D 卷积元语基准实现)

CPU 上的其他 3D CNN 元语操作

除了 3D 卷积元语以外,还必须确保所有相关 3D元语的速度。得益于英特尔® MKL 和英特尔® MKL-DNN,我们实现了一系列CPU上的快速 3D 元语操作,包括 3D 批量归一化、3D 反卷积、3D 池化、3D SoftMax 损失、3D 交叉熵损失、3D smooth L1 损失、3D concat 等。这些成果已经发布到Extended-Caffe [10],并成为天池大赛软件堆栈的核心。图 5 显示了整体模型的训练和推理速度的提升。

图 5.整体效率提升
(优化前、后对比)。
训练耗时按照一次迭代来测量(即处理一个裁剪区的时间),而推理耗时按照一整张CT 图像来测量(即处理一幅完整的CT 图像中所有裁剪区的时间)。

实验结果和定量分析

模型训练

我们采用随机梯度下降(SGD)和逐步学习率策略来训练模型。所有网络参数都进行随机初始化,初始学习率设为 0.01。模型训练 100 次,并分别将第 50 次和第 80 次的学习率缩小 10 倍。图 6 记录了我们训练 LUNA'16 数据集(例如子集 0)时的损失趋势。可以看到,在第 80 至第 100 轮训练期间,模型在该数据集上成功收敛。

图 6.模型训练示例(LUNA’16 subset 0)

检测性能评估

天池、LUNA'16 等知名大赛都采用 FROC(自由响应接收器操作特性)[11] 方法来评估模型的检测性能。如果某一候选区(即模型最终输出的立方体区域框)的中心点与实际结节的中心点的距离比结节半径短,那么该候选区将被视作真阳性。然后,计算 FROC 得分以综合评价该模型的敏感性与假阳性误差。

分辨率的影响

图 7 显示了 3D CNN 模型的不同 FROC 得分与用于模型训练的输入数据的不同分辨率之间的关系。可以看到,分辨率较高时训练模型的 FROC 得分较高。也就是说,实验证明,高分辨率有助于提高模型的检测性能。

图 7.分辨率越高,FROC 得分越高

由于人类放射科医生能够轻易检测到大结节,而很难检测到小结节,因此急需 AI 辅助检测方案具备高效检测小结节的能力。图 8 比较了我们的模型在不同分辨率下检测不同尺寸的结节的效果。可以看到,高分辨率可显著提高小结节的检测性能。

图 8.高分辨率可提高检测小结节的精确度

内存使用分析

在不同分辨率下训练模型的过程中,我们分析了内存使用情况,并对比了 CPU 平台和 GPGPU 平台的行为。图 9 (a) 和 (b) 记录了训练批量大小(batch size)分别等于 1 和 4 的情况。当批量大小等于 1 时,具备 12GB 内存的现代 GPGPU 仅支持最高 128 x 128 x 128 的分辨率,而包含 384GB 内存 CPU 平台可轻松支持高达 448 x 448 x 448 的分辨率。而当批量大小等于 4 时,该GPGPU 的情况则更糟 — 仅支持最高 96 x 96 x 96 的分辨率,而 上述CPU 平台可轻松支持高达 256 x 256 x 256 的分辨率。


(a) 批量大小=1


(b) 批量大小=4

图 9. 内存使用与不同分辨率

由于现代 CPU 服务器拥有数 TB 内存容量,而且考虑到英特尔即将推出Apache Pass 技术,因此可以说, CPU 平台能提供近乎无限的灵活性,用以支持医学影像分析领域的模型设计人员探索分辨率极高的解决方案,从而获取最佳检测性能。

总结

由阿里云、英特尔和零氪科技联合主办的天池医疗 AI 大赛,其环境是一套完全基于英特尔深度学习软硬件堆栈构建而成的AI云计算服务。本文从天池大赛出发,以定量实验的方式阐述了我们在医学 AI 领域所获得的重要结论。首先,我们自主研发了一个可反映当前最新肺结节检测水平的 3D CNN 模型。接下来,我们通过定量方式证明了,用高分辨率数据训练的模型可显著提升检测性能,特别是检测小结节的精确度。同时,高分辨率模型必须消耗更多内存。我们对比了 GPGPU 和 CPU,并证明 CPU 平台由于能够提供较大的内存容量,有助于医疗 AI 设计人员探索更高分辨率的解决方案,进而实现最佳的检测性能。工程上,通过实现 Extended-Caffe 框架作为天池大赛的核心软件,以此为例,证明了 CPU 架构可以支持快速 3D CNN 运算,从而保证了医疗影像分析领域的开发者在CPU上真正高效且快速的研发3D CNN 模型。

参考

  1.  https://tianchi.aliyun.com/getStart/introduction.htm?raceId=231601
  2. Bush I., Lung nodule detection and classification. Technical report, Stanford Computer Science, 2016.
  3. https://luna16.grand-challenge.org/home/
  4. Girshick, R. Fast R-CNN. Computer Science, 2015.
  5. Ronneberger, O., Fischer, P., and Brox, T. U-Net: Convolutional Networks for Biomedical Image Segmentation, Medical Image Computing and Computer-Assisted Intervention — MICCAI 2015. Springer International Publishing, 2015:234‒241.
  6. https://github.com/lfz/DSB2017
  7. https://software.intel.com/zh-cn/articles/intel-mkl-dnn-part-1-library-overview-and-installation
  8. https://software.intel.com/zh-cn/mkl
  9. https://luna16.grand-challenge.org/results/
  10. https://github.com/extendedcaffe/extended-caffe
  11. http://devchakraborty.com/Receiver%20operating%20characteristic.pdf
有关编译器优化的更完整信息,请参阅优化通知