NAS-Unet: Neural Architecture Search for Medical Image Segmentation

TLDR

Neural architecture search has advanced image‑classification accuracy, and recent efforts extend it to segmentation, yet existing studies target natural scenes rather than medical images. This work introduces NAS‑Unet, which automatically discovers DownSC and UpSC cells for medical image segmentation by searching a U‑like backbone. The method trains DownSC and UpSC cells concurrently using a differentiable search strategy within a U‑like backbone. On Promise12, Chaos, and ultrasound nerve datasets, the searched architecture achieves superior segmentation accuracy with only ~0.8 M parameters, outperforming U‑net and its variants without pretraining.

Abstract

Neural architecture search (NAS) has significant progress in improving the accuracy of image classification. Recently, some works attempt to extend NAS to image segmentation which shows preliminary feasibility. However, all of them focus on searching architecture for semantic segmentation in natural scenes. In this paper, we design three types of primitive operation set on search space to automatically find two cell architecture DownSC and UpSC for semantic image segmentation especially medical image segmentation. Inspired by the U-net architecture and its variants successfully applied to various medical image segmentation, we propose NAS-Unet which is stacked by the same number of DownSC and UpSC on a U-like backbone network. The architectures of DownSC and UpSC updated simultaneously by a differential architecture strategy during the search stage. We demonstrate the good segmentation results of the proposed method on Promise12, Chaos, and ultrasound nerve datasets, which collected by magnetic resonance imaging, computed tomography, and ultrasound, respectively. Without any pretraining, our architecture searched on PASCAL VOC2012, attains better performances and much fewer parameters (about 0.8M) than U-net and one of its variants when evaluated on the above three types of medical image datasets.