[论文解读] Resource Optimized Neural Architecture Search for 3D Medical Image Segmentation

Woong Bae, Seungho Lee+4 more|arXiv (Cornell University)|Sep 2, 2019

Advanced Neural Network Applications参考文献 10被引用 5

一句话总结

该论文提出 RONASMIS，一种面向 3D 医疗图像分割的资源优化神经架构搜索框架，仅使用一张 RTX 2080Ti（10.8GB GPU 内存）并在 1.39 天内完成 1GB 数据集的训练，即实现了最先进性能。通过聚焦于宏搜索、参数共享、采用逐元素相加的跳跃连接并避免重新训练，该方法在无需数据增强、模型集成或推理时增强的情况下，优于手动设计的网络及现有 NAS 方法。

ABSTRACT

Neural Architecture Search (NAS), a framework which automates the task of designing neural networks, has recently been actively studied in the field of deep learning. However, there are only a few NAS methods suitable for 3D medical image segmentation. Medical 3D images are generally very large; thus it is difficult to apply previous NAS methods due to their GPU computational burden and long training time. We propose the resource-optimized neural architecture search method which can be applied to 3D medical segmentation tasks in a short training time (1.39 days for 1GB dataset) using a small amount of computation power (one RTX 2080Ti, 10.8GB GPU memory). Excellent performance can also be achieved without retraining(fine-tuning) which is essential in most NAS methods. These advantages can be achieved by using a reinforcement learning-based controller with parameter sharing and focusing on the optimal search space configuration of macro search rather than micro search. Our experiments demonstrate that the proposed NAS method outperforms manually designed networks with state-of-the-art performance in 3D medical image segmentation.

研究动机与目标

• 解决现有 NAS 方法在 3D 医疗图像分割中计算成本高和训练时间长的问题。
• 降低 3D 医疗影像中 NAS 的 GPU 显存占用和训练时间，因数据集大且具有各向异性特征。
• 开发一种 NAS 框架，其性能优于人工设计的架构，且无需重新训练或大量数据增强。
• 实现在计算资源有限的临床环境中 NAS 的实际部署。
• 证明在 3D 医疗分割任务中，通过战略性架构选择的宏搜索可超越微搜索。

提出的方法

• 该框架采用基于强化学习的控制器并结合参数共享，在宏搜索空间中高效搜索，从而降低训练成本和显存使用。
• 搜索空间包含输入块大小、下采样操作、池化类型、空洞率和跳跃连接点，所有选择均以在最小资源消耗下最大化性能为目标。
• 使用逐元素相加代替拼接实现跳跃连接，以减少参数量和 GPU 显存消耗。
• 将深度可分离卷积替换为标准 3D 卷积，以提升显存效率和性能。
• 应用丢弃路径正则化以防止子网络在搜索过程中过拟合，同时用实例归一化替代批量归一化以稳定训练。
• 控制器训练 150–500 个周期，每个训练周期内子网络训练 3 个周期，奖励基于验证集 Dice 评分。

实验结果

研究问题

• RQ1 能否设计一种适用于 3D 医疗图像分割的 NAS 框架，使其在有限 GPU 资源下高效运行？
• RQ2 在 3D 医疗图像分割任务中，通过战略性架构选择的宏搜索是否优于微搜索？
• RQ3 能否实现一种 NAS 方法，在无需重新训练、数据增强或模型集成的情况下达到最先进性能？
• RQ4 参数共享与逐元素跳跃连接在降低显存使用的同时，对性能保持效果如何？
• RQ5 基于控制器的 NAS 方法能否在不依赖预训练模型或复杂数据流水线的情况下，实现稳定且一致的架构搜索？

主要发现

• RONASMIS 在心脏分割任务中达到 92.72 的平均 Dice 评分，超过 nnUnet（92.70）和 SCNAS（89.99），展现出最先进性能。
• 在前列腺分割任务中，该方法取得 75.71 的 Dice 评分，优于 nnUnet（74.54）和 SCNAS（65.30），且未使用推理时增强或模型集成。
• 在单张 RTX 2080Ti（10.8GB GPU 显存）上，1GB 数据集的训练时间缩短至 1.39 天，显著低于以往 NAS 方法。
• 控制器在超过 80% 的时间里建议相同或高度相似的架构，表明其搜索行为稳定且一致，熵值随时间下降，奖励持续上升。
• 该方法在无需数据增强、模型集成、推理时增强或后处理的情况下实现卓越性能，而大多数最先进基线方法均依赖这些技术。
• 该框架成功将深度可分离卷积替换为标准 3D 卷积，并采用实例归一化以稳定训练，从而实现在 3D 医疗图像上高效且稳定的 NAS。