SIDN-NAS: Scalable Iterative Dense Network with Neural Architecture Search Optimization for Medical Image Segmentation
Jianjun Zhou, Junying Chen · 2024
Medical image segmentation faces problems such as the unbalanced foreground and background of medical images and the limited dataset size. Meanwhile, Neural Architecture Search (NAS) methods have shown the great application potential and value in medical image segmentation tasks. Compared to the manually-designed medical image segmentation models, the NAS methods can improve the model accuracy, but the complexity and the computing cost of the searched medical image segmentation model is relatively high. In this work, we propose a high-precision but low-complexity scalable medical image segmentation model (Scalable Iterative Dense Network, SIDN) with NAS optimization. The proposed SIDN model consists of scalable iterative U-shape encoder/decoder stacks, uses fewer convolutional kernels in each encoder or decoder and performs dense multi-scale residual semantic connections between encoders and decoders, which can make up for the semantic differences between different levels. It then uses a gradient-based NAS algorithm to further reduce the model complexity and improve the segmentation accuracy. The SIDN model training and NAS optimization only require 1 GPU hour, and the segmentation accuracy has been improved by 7% to 10% on multiple medical image datasets. More importantly, the number of model parameters has been reduced by 31% to 88%, and the amount of calculation has been reduced by 32% to 44 %, greatly improving the model inference speed. Our source code is available at: https://github.com/Bob5090/SIDN-NAS.