NG-NAS: Node growth neural architecture search for 3D medical image segmentation

Comput Med Imaging Graph. 2023 Sep:108:102268. doi: 10.1016/j.compmedimag.2023.102268. Epub 2023 Jun 16.

Abstract

Neural architecture search (NAS) has been applied to design proper 3D networks for medical image segmentation. In order to reduce the computation cost in NAS, researchers tend to adopt weight sharing mechanism to search architectures in a supernet. However, recent studies state that the searched architecture rankings may not be accurate with weight sharing mechanism because the training situations are inconsistent between the searching and training phases. In addition, some NAS algorithms design inflexible supernets that only search operators in a pre-defined backbone and ignore the importance of network topology, which limits the performance of searched architecture. To avoid weight sharing mechanism which may lead to inaccurate results and to comprehensively search network topology and operators, we propose a novel NAS algorithm called NG-NAS. Following the previous studies, we consider the segmentation network as a U-shape structure composed of a set of nodes. Instead of searching from the supernet with a limited search space, our NG-NAS starts from a simple architecture with only 5 nodes, and greedily grows the best candidate node until meeting the constraint. We design 2 kinds of node generations to form various network topological structures and prepare 4 candidate operators for each node. To efficiently evaluate candidate node generations, we use NAS without training strategies. We evaluate our method on several public 3D medical image segmentation benchmarks and achieve state-of-the-art performance, demonstrating the effectiveness of the searched architecture and our NG-NAS. Concretely, our method achieves an average Dice score of 85.11 on MSD liver, 65.70 on MSD brain, and 87.59 in BTCV, which performs much better than the previous SOTA methods.

Keywords: 3D medical image segmentation; AutoML; Deep learning; Neural architecture search; Node growth.

Publication types

  • Research Support, Non-U.S. Gov't

MeSH terms

  • Algorithms*
  • Benchmarking*
  • Brain / diagnostic imaging
  • Image Processing, Computer-Assisted
  • Liver

Substances

  • nas