[论文解读] High Resolution Medical Image Analysis with Spatial Partitioning

Le Hou, Youlong Cheng+7 more|arXiv (Cornell University)|Sep 6, 2019

Medical Image Segmentation Techniques参考文献 18被引用 29

一句话总结

该论文提出了一种在Mesh-TensorFlow中的空间划分框架，实现了对高达512×512×512体素的高分辨率医学图像的3D U-Net模型端到端训练，通过全息交换（halo exchange）保持分割块之间的卷积连接性。据我们所知，这是首个直接在完整分辨率CT扫描上训练此类模型而无需下采样或数据裁剪的方法，在训练时间仅增加5%的情况下，于LiTS基准上实现了最先进（SOTA）的Dice分数。

ABSTRACT

Medical images such as 3D computerized tomography (CT) scans and pathology images, have hundreds of millions or billions of voxels/pixels. It is infeasible to train CNN models directly on such high resolution images, because neural activations of a single image do not fit in the memory of a single GPU/TPU, and naive data and model parallelism approaches do not work. Existing image analysis approaches alleviate this problem by cropping or down-sampling input images, which leads to complicated implementation and sub-optimal performance due to information loss. In this paper, we implement spatial partitioning, which internally distributes the input and output of convolutional layers across GPUs/TPUs. Our implementation is based on the Mesh-TensorFlow framework and the computation distribution is transparent to end users. With this technique, we train a 3D Unet on up to 512 by 512 by 512 resolution data. To the best of our knowledge, this is the first work for handling such high resolution images end-to-end.

研究动机与目标

• 解决由于GPU/TPU显存限制，导致在高分辨率医学图像（如包含10⁸–10⁹个体素的3D CT扫描）上直接训练CNN的不可行性。
• 克服现有方法（如裁剪、下采样或粗到细方案）导致的信息损失和复杂实现的局限性。
• 实现在多个GPU/TPU上对全分辨率3D医学图像分析的高效、透明的模型与数据并行。
• 开发一个基于Mesh-TensorFlow的可扩展、开源框架，支持在高分辨率数据上进行TPU和GPU训练。
• 通过引入一种基于合成的新数据增强方法，提升肝脏肿瘤分割的泛化能力。

提出的方法

• 空间划分将高分辨率3D图像划分为非重叠的块，并分布到多个GPU/TPU上。
• 在每个卷积层之前执行全息交换，相邻设备交换块边缘（一半卷积核大小）以保持空间上下文。
• 该框架基于Mesh-TensorFlow构建，可实现自动的数据与模型并行，用户干预极少。
• 一种基于合成的数据增强方法生成具有真实强度与形状变化的合成肿瘤，以提升模型鲁棒性。
• 该方法支持在无需下采样或基于块的推理下，对512×512×512分辨率数据进行端到端3D U-Net训练。
• 训练采用混合精度（半精度）浮点运算，优化器为Adafactor，损失函数结合Dice与交叉熵损失。

实验结果

研究问题

• RQ1 能否在不进行下采样或裁剪的情况下，对完整分辨率512×512×512 CT扫描实现3D U-Net模型的端到端训练？
• RQ2 结合全息交换的空间划分是否能实现高效、可扩展的高分辨率医学图像训练，并保持极低的计算开销？
• RQ3 基于合成的数据增强方法是否能提升3D CT扫描中肝脏肿瘤分割的泛化能力与性能？
• RQ4 在所提框架下，LiTS基准上的Dice分数如何随输入分辨率的提升而变化？
• RQ5 在分布式TPU环境中，空间划分与全息交换实际引入的训练时间开销是多少？

主要发现

• 所提框架成功实现了在512×512×512分辨率CT扫描上的3D U-Net模型训练，在LiTS验证集上达到平均每个病例Dice分数为0.4547 ± 0.0475。
• 在512×512×512分辨率下，Dice全局分数达到0.7180 ± 0.0446，表明在大体积分割任务中表现优异。
• 由于划分、重塑和全息交换操作，仅引入约5%的额外训练时间开销。
• 若不采用所提数据增强方法，每个病例的Dice分数至少下降10%，证明其在模型泛化中的关键作用。
• 更高分辨率输入（如512³）始终优于较低分辨率（如64³、128³），证实全分辨率学习的优势。
• 该框架支持TPU与GPU训练，并可在不修改网络架构的前提下，透明地实现计算在设备间的分布。