SegNAS3D: Network Architecture Search with Derivative-Free Global Optimization for 3D Image Segmentation

TL;DR

SegNAS3D introduces a derivative-free global optimization framework for automated 3D medical image segmentation architecture search, achieving high accuracy with smaller models in less than three days.

Contribution

This work presents a novel architecture search method specifically for 3D image segmentation, addressing the gap in automated design for medical imaging tasks.

Findings

Achieved an average Dice coefficient of 82% on 3D brain MRI segmentation.

Completed architecture search in less than three days on three GPUs.

Produced smaller architectures compared to state-of-the-art manual designs.

Abstract

Deep learning has largely reduced the need for manual feature selection in image segmentation. Nevertheless, network architecture optimization and hyperparameter tuning are mostly manual and time consuming. Although there are increasing research efforts on network architecture search in computer vision, most works concentrate on image classification but not segmentation, and there are very limited efforts on medical image segmentation especially in 3D. To remedy this, here we propose a framework, SegNAS3D, for network architecture search of 3D image segmentation. In this framework, a network architecture comprises interconnected building blocks that consist of operations such as convolution and skip connection. By representing the block structure as a learnable directed acyclic graph, hyperparameters such as the number of feature channels and the option of using deep supervision can be learned together through derivative-free global optimization. Experiments on 43 3D brain magnetic resonance images with 19 structures achieved an average Dice coefficient of 82%. Each architecture search required less than three days on three GPUs and produced architectures that were much smaller than the state-of-the-art manually created architectures.