EMONAS-Net: Efficient multiobjective neural architecture search using surrogate-assisted evolutionary algorithm for 3D medical image segmentation

Abstract

Deep learning plays a critical role in medical image segmentation. Nevertheless, manually designing a neural network for a specific segmentation problem is a very difficult and time-consuming task due to the massive hyperparameter search space, long training time and large volumetric data. Therefore, most designed networks are highly complex, task specific and over-parametrized. Recently, multiobjective neural architecture search (NAS) methods have been proposed to automate the design of accurate and efficient segmentation architectures. However, they only search for either the micro- or macro-structure of the architecture, do not use the information produced during the optimization process to increase the efficiency of the search, or do not consider the volumetric nature of medical images. In this work, we present EMONAS-Net, an Efficient MultiObjective NAS framework for 3D medical image segmentation that optimizes both the segmentation accuracy and size of the network. EMONAS-Net has two key components, a novel search space that considers the configuration of the micro- and macro-structure of the architecture and a Surrogate-assisted Multiobjective Evolutionary based Algorithm (SaMEA algorithm) that efficiently searches for the best hyperparameter values. The SaMEA algorithm uses the information collected during the initial generations of the evolutionary process to identify the most promising subproblems and select the best performing hyperparameter values during mutation to improve the convergence speed. Furthermore, a Random Forest surrogate model is incorporated to accelerate the fitness evaluation of the candidate architectures. EMONAS-Net is tested on the tasks of prostate segmentation from the MICCAI PROMISE12 challenge, hippocampus segmentation from the Medical Segmentation Decathlon challenge, and cardiac segmentation from the MICCAI ACDC challenge. In all the benchmarks, the proposed framework finds architectures that perform better or comparable with competing state-of-the-art NAS methods while being considerably smaller and reducing the architecture search time by more than 50%.