Multi-Objective Once-for-All Neural Architecture Search for Medical Image Segmentation

Deep learning is the mainstream method for medical image segmentation, and neural architecture search (NAS) has also been developed for this task. However, existing NAS methods remain limited in their ability to search for high-performance yet lightweight network architectures due to the high computational cost of NAS and the low fidelity of performance evaluation during the search process. In this paper, we propose a novel once-for-all NAS method for medical image segmentation to address these challenges. A carefully designed search space (supernet) incorporating key components of U-shape networks is constructed specifically for medical image segmentation. An effective and efficient hybrid two-stage supernet training scheme is then designed to enhance supernet training while maintaining a balance between performance and computational cost. A multi-objective evolutionary algorithm is leveraged to search for sets of network architectures, which produces high-performing architectures with varying computational complexities, optimized under multiple objectives. We conduct experiments on six widely used medical image segmentation datasets. Compared with existing methods, the proposed method achieves state-of-the-art performance on all six datasets. The searched architectures exhibit an excellent trade-off between performance and computational complexity, which is attributed to the effective multi-objective search. Our source codes are available at https://github.com/jiahongwei21-lang/MOOFA4MIS.