Auto-nnU-Net: Towards Automated Medical Image Segmentation

TL;DR

Auto-nnU-Net addresses the MIS model selection challenge across diverse tasks by integrating hyperparameter optimization, neural architecture search, and hierarchical NAS into the nnU-Net framework, complemented by Regularized PriorBand to balance segmentation accuracy with training runtime. The approach is validated on the Medical Segmentation Decathlon, where it improves performance on the majority of datasets and maintains practical resource use, outperforming several baselines. The work provides insights into hyperparameter importance, cross-dataset transferability of incumbents, and a scalable AutoML workflow for MIS, with public code to enable reproducibility. This framework holds practical significance for deploying accurate MIS systems in resource-constrained clinical environments.

Abstract

Medical Image Segmentation (MIS) includes diverse tasks, from bone to organ segmentation, each with its own challenges in finding the best segmentation model. The state-of-the-art AutoML-related MIS-framework nnU-Net automates many aspects of model configuration but remains constrained by fixed hyperparameters and heuristic design choices. As a full-AutoML framework for MIS, we propose Auto-nnU-Net, a novel nnU-Net variant enabling hyperparameter optimization (HPO), neural architecture search (NAS), and hierarchical NAS (HNAS). Additionally, we propose Regularized PriorBand to balance model accuracy with the computational resources required for training, addressing the resource constraints often faced in real-world medical settings that limit the feasibility of extensive training procedures. We evaluate our approach across diverse MIS datasets from the well-established Medical Segmentation Decathlon, analyzing the impact of AutoML techniques on segmentation performance, computational efficiency, and model design choices. The results demonstrate that our AutoML approach substantially improves the segmentation performance of nnU-Net on 6 out of 10 datasets and is on par on the other datasets while maintaining practical resource requirements. Our code is available at https://github.com/automl/AutoNNUnet.

Figures (4)

• Figure 1: Overview of the Auto-nnU-Net framework: Given a hyperparameter configuration $\lambda$, architecture $A$, and dataset fingerprint, the AutoExperimentPlanner and CFGUNet generate nnU-Net training plans and model architecture, respectively. The AutoNNUNetTrainer then trains the selected model, providing runtime and validation score as objectives to the PriorBand optimizer. For details, see Appendix \ref{['sec:appendix:approach:autonnunet']}. For more details, see Appendix \ref{['sec:appendix:approach:autonnunet']}.
• Figure 2: Incumbent performance of nnU-Net, Auto-nnU-Net, and Auto-nnU-Net ablations over time. Detailed results for each dataset are stated in Appendix \ref{['sec:appendix:results']}. Final validation DSCs are stated in Table \ref{['tab:results:dsc_overview']} in Appendix \ref{['sec:appendix:results']}. We exclude MedSAM2 as it fails to achieve the performance of nnU-Net on nine out of ten datasets (see Table \ref{['tab:results:dsc_overview']}).
• Figure 3: Qualitative segmentation results for D01 . Columns show the input image, ground truth mask, and method predictions, with colors denoting foreground classes. Numbers below method names indicate DSC scores [%] for this example. Each row shows a slice of the 3D volume along one axis. As the 4D volume is an mp-MRI scan, the first parameter setting is used to extract a 3D volume. Additional results are in Appendix \ref{['sec:appendix:results']}.
• Figure 4: Pareto fronts of Auto-nnU-Net and HPO+HNAS for D03 (left) and D04 (right) compared to the baselines and HPO results. Additional results are stated in Appendix \ref{['sec:appendix:results']}.