[论文解读] V-NAS: Neural Architecture Search for Volumetric Medical Image Segmentation

Zhuotun Zhu, Chenxi Liu+3 more|arXiv (Cornell University)|Jun 6, 2019

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

一句话总结

该论文提出V-NAS，一种可微神经架构搜索框架，可自动学习用于体素医学图像分割的2D、3D及伪3D（P3D）卷积的最佳组合。通过将架构搜索建模为连续可微的优化问题，V-NAS在三个公开数据集上的正常与异常器官分割任务中均实现了最先进性能，且所搜索到的架构在不同数据集间具有良好泛化能力。

ABSTRACT

Deep learning algorithms, in particular 2D and 3D fully convolutional neural networks (FCNs), have rapidly become the mainstream methodology for volumetric medical image segmentation. However, 2D convolutions cannot fully leverage the rich spatial information along the third axis, while 3D convolutions suffer from the demanding computation and high GPU memory consumption. In this paper, we propose to automatically search the network architecture tailoring to volumetric medical image segmentation problem. Concretely, we formulate the structure learning as differentiable neural architecture search, and let the network itself choose between 2D, 3D or Pseudo-3D (P3D) convolutions at each layer. We evaluate our method on 3 public datasets, i.e., the NIH Pancreas dataset, the Lung and Pancreas dataset from the Medical Segmentation Decathlon (MSD) Challenge. Our method, named V-NAS, consistently outperforms other state-of-the-arts on the segmentation task of both normal organ (NIH Pancreas) and abnormal organs (MSD Lung tumors and MSD Pancreas tumors), which shows the power of chosen architecture. Moreover, the searched architecture on one dataset can be well generalized to other datasets, which demonstrates the robustness and practical use of our proposed method.

研究动机与目标

• 为解决固定2D、3D或P3D卷积网络在体素医学图像分割中的局限性，这些网络要么忽略3D空间上下文，要么计算成本过高。
• 开发一种自动化方法，学习最优的架构配置——在每一层动态选择2D、3D或P3D卷积，以适配体素医学数据的特定特征。
• 评估所搜索架构在不同医学影像数据集间的泛化能力，特别是对正常与异常器官分割任务的适用性。
• 证明可微神经架构搜索在分割精度方面可超越人工设计的架构，同时保持计算效率。

提出的方法

• 将体素分割的神经架构搜索建模为可微的连续优化问题，支持架构参数的端到端训练。
• 引入一个搜索空间，使每一层可动态选择2D、3D或伪3D（P3D）卷积，其中P3D通过两个连续的2D卷积（k×k×1后接1×1×k）替代3D核。
• 通过类似DARTS的可微架构搜索方法对架构搜索空间进行连续松弛，实现对每一层卷积类型选择的基于梯度的优化。
• 采用单阶段可微搜索流程进行网络训练，无需预训练，直接联合优化架构与权重。
• 采用单路径搜索策略，使用一组共享权重，高效探索多个数据集的架构空间。
• 在NIH胰腺、MSD肺部肿瘤和MSD胰腺肿瘤数据集的保留测试集上评估最终搜索到的架构，以评估性能与泛化能力。

实验结果

研究问题

• RQ1 可微神经架构搜索框架能否自动发现用于体素医学图像分割的2D、3D及伪3D卷积的最佳组合？
• RQ2 所提出的V-NAS方法在分割精度方面是否优于固定卷积类型（如纯2D、3D或P3D）的人工设计架构？
• RQ3 在某一数据集上搜索到的架构能否有效泛化到具有不同解剖结构与病理特征的其他医学影像数据集？
• RQ4 V-NAS在正常器官与异常肿瘤分割任务上的性能与当前最先进方法相比如何？

主要发现

• 在NIH胰腺数据集上，V-NAS实现了85.15%的平均Dice相似系数（DSC），优于所有人工配置的架构，包括最佳人工配置（P3D/P3D，DSC为84.75%）。
• 在MSD肺部肿瘤数据集上，V-NAS实现了55.27%的平均DSC，较最佳人工配置（3D/3D，DSC为53.74%）高出1.53%。
• 在MSD胰腺肿瘤数据集上，V-NAS实现了最先进性能，表明其在正常与异常器官分割任务中均具有一致优势。
• 在某一数据集上搜索到的架构能良好泛化至其他数据集，表明其具备鲁棒性与在真实临床部署中的实用价值。
• 消融实验证实，通过NAS实现的卷积类型自动选择优于固定的人工配置，凸显了架构灵活性在应对多样化解剖与病理变异中的重要性。
• 该方法计算效率高，无需预训练，支持架构与权重的直接端到端联合优化。