[论文解读] Fast and Accurate 3D Medical Image Segmentation with Data-swapping Method

Haruki Imai, Samuel Matzek+3 more|arXiv (Cornell University)|Dec 19, 2018

Advanced Neural Network Applications被引用 8

一句话总结

本文提出了一种数据交换方法，通过在GPU和CPU内存之间动态移动中间特征，实现了在不使用分块处理的情况下对完整尺寸的192×192×192医学图像进行3D U-Net模型训练。该方法将通信开销降低了17.1%，全肿瘤和肿瘤核心的Dice分数分别提高了4.48%和5.32%，训练速度提升了3.53倍（从164小时缩短至47小时）。

ABSTRACT

Deep neural network models used for medical image segmentation are large because they are trained with high-resolution three-dimensional (3D) images. Graphics processing units (GPUs) are widely used to accelerate the trainings. However, the memory on a GPU is not large enough to train the models. A popular approach to tackling this problem is patch-based method, which divides a large image into small patches and trains the models with these small patches. However, this method would degrade the segmentation quality if a target object spans multiple patches. In this paper, we propose a novel approach for 3D medical image segmentation that utilizes the data-swapping, which swaps out intermediate data from GPU memory to CPU memory to enlarge the effective GPU memory size, for training high-resolution 3D medical images without patching. We carefully tuned parameters in the data-swapping method to obtain the best training performance for 3D U-Net, a widely used deep neural network model for medical image segmentation. We applied our tuning to train 3D U-Net with full-size images of 192 x 192 x 192 voxels in brain tumor dataset. As a result, communication overhead, which is the most important issue, was reduced by 17.1%. Compared with the patch-based method for patches of 128 x 128 x 128 voxels, our training for full-size images achieved improvement on the mean Dice score by 4.48% and 5.32 % for detecting whole tumor sub-region and tumor core sub-region, respectively. The total training time was reduced from 164 hours to 47 hours, resulting in 3.53 times of acceleration.

研究动机与目标

• 解决在高分辨率医学图像上训练深度3D U-Net模型时GPU内存受限的问题。
• 消除基于分块训练方法导致的物体在分块间断裂造成的质量下降问题。
• 开发一种高效管理GPU与CPU内存之间中间激活值的数据交换技术，以模拟更大的有效GPU内存。
• 在不损害解剖结构空间连续性的情况下，实现更快、更精确的3D医学图像分割。

提出的方法

• 该方法在训练过程中将中间特征图从GPU内存卸载到CPU内存，以减轻GPU内存压力。
• 该方法基于内存使用情况和计算依赖关系，动态调度GPU与CPU之间的数据传输。
• 该方法专门针对3D U-Net进行调优，通过优化数据传输频率和大小，最大限度减少通信开销。
• 该框架支持在不使用分块处理的情况下，对完整的192×192×192体素图像进行端到端训练，从而保留空间上下文信息。
• 该系统采用自定义的内存管理策略，优先处理频繁访问的特征，并最小化冗余传输。
• 该方法在脑肿瘤分割数据集上通过Dice分数和训练时间等标准指标进行了评估。

实验结果

研究问题

• RQ1 数据交换是否能够在不使用分块处理的情况下，实现对全尺寸3D医学图像的3D U-Net训练，同时保持或提升分割精度？
• RQ2 与标准的基于分块的训练相比，数据交换的通信开销在训练效率方面表现如何？
• RQ3 避免分块处理在小尺寸或碎片化肿瘤区域的分割性能提升程度如何？
• RQ4 为最小化训练时间并最大化模型精度，数据交换参数的最佳配置是什么？
• RQ5 与基于分块的替代方案相比，该数据交换方法是否能够实现更快的收敛速度和更好的泛化能力？

主要发现

• 与基线方法相比，数据交换方法将通信开销降低了17.1%，显著提升了训练效率。
• 在完整192×192×192图像上进行训练，全肿瘤子区域的平均Dice分数相比基于分块的训练提高了4.48%。
• 该方法在肿瘤核心子区域的Dice分数提高了5.32%，表明对关键区域的定位能力更强。
• 总训练时间从164小时减少到47小时，训练速度提升了3.53倍。
• 该方法在准确性和速度上均优于基于分块的训练，证明了保留完整体素上下文信息的优势。
• 结果证实，数据交换技术能够在不引入分块相关伪影和性能损失的前提下，实现高分辨率3D分割。