[论文解读] LAMP: Large Deep Nets with Automated Model Parallelism for Image Segmentation

Wentao Zhu, Can Zhao+4 more|arXiv (Cornell University)|Jun 22, 2020

Advanced Neural Network Applications参考文献 29被引用 6

一句话总结

LAMP 将自动模型并行化引入大 3D U-Net 和 SEU-Net 架构的训练中，采用完整图像输入，显著提升了分割精度和推理速度。通过基于 GPipe 的并行 U-Net 设计，该方法在无需修改模型架构的情况下，在医学图像分割任务中实现了最先进性能。

ABSTRACT

Deep Learning (DL) models are becoming larger, because the increase in model size might offer significant accuracy gain. To enable the training of large deep networks, data parallelism and model parallelism are two well-known approaches for parallel training. However, data parallelism does not help reduce memory footprint per device. In this work, we introduce Large deep 3D ConvNets with Automated Model Parallelism (LAMP) and investigate the impact of both input's and deep 3D ConvNets' size on segmentation accuracy. Through automated model parallelism, it is feasible to train large deep 3D ConvNets with a large input patch, even the whole image. Extensive experiments demonstrate that, facilitated by the automated model parallelism, the segmentation accuracy can be improved through increasing model size and input context size, and large input yields significant inference speedup compared with sliding window of small patches in the inference. Code is available\footnote{https://monai.io/research/lamp-automated-model-parallelism}.

研究动机与目标

• 研究模型尺寸和输入上下文尺寸增大对 3D 医学图像分割精度的影响。
• 解决 GPU 显存限制问题，该问题阻碍了在医学影像中训练大 3D ConvNets 和完整图像输入。
• 开发一种与标准 U-Net 架构兼容的自动模型并行化框架，该架构在医学图像分析中被广泛使用。
• 通过使用大感受野的单次完整图像推理替代滑动窗口推理，减少推理时间。
• 验证课程训练和自适应加权损失在处理医学图像分割中类别不平衡问题的有效性。

提出的方法

• 基于 GPipe 的流水线并行化设计并行 U-Net 架构，将大 3D ConvNet 层分布到多个 GPU 上。
• 实现自动模型并行化，将网络跨设备划分，无需手动分层或修改架构。
• 采用课程训练策略，按顺序在小、中、大输入块上进行训练，以稳定大模型的学习过程。
• 应用结合 Dice 损失和焦点损失的自适应加权损失，以解决小器官或病灶医学图像中的类别不平衡问题。
• 使用 RMSProp 优化器配合学习率调度和批量归一化，以稳定大模型的训练过程。
• 通过使用大输入块（最大 512×512×704）实现完整图像推理，消除对滑动窗口推理的需求。

实验结果

研究问题

• RQ1 增大 3D ConvNets 尺寸对医学图像数据集分割精度有何影响？
• RQ2 与滑动窗口推理相比，使用大输入上下文（完整图像或大块）对推理速度和精度有何影响？
• RQ3 自动模型并行化能否在不修改架构的前提下，使标准医学影像数据集上的大 3D U-Net 和 SEU-Net 模型得以训练？
• RQ4 通过逐步增大输入尺寸的课程训练策略，如何改善大模型的收敛性和性能？
• RQ5 结合自适应加权损失、大模型和大输入，在小结构分割中能在多大程度上缓解类别不平衡问题？

主要发现

• 采用完整图像输入的 SEU-Net-128 模型在性能上优于 AnatomyNet，表明大模型和大上下文可显著提升精度。
• 使用完整图像输入将单张图像的推理时间缩短至 1.52 秒（U-Net-32 在 256³ 输入下），相比使用 64³ 块的滑动窗口推理提速 5.8 倍。
• 大模型（如 U-Net-128）配合大输入在肝脏分割中取得 90.99% 的 Dice 分数，在肿瘤分割中取得 56.48% 的 Dice 分数，显著优于较小配置。
• 对于最大输入（256³），U-Net-128 的推理时间缩短至 4.39 秒，表明大输入可通过减少前向传播次数实现更快推理。
• 课程训练策略实现了大模型的稳定训练，且在逐步增大块尺寸上训练时获得了更高精度。
• 自动模型并行化框架成功训练了每块最多 128 个滤波器、输入尺寸最大达 512×512×704 的模型，突破了 GPU 显存限制。