Source-linked AI summary
nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation
Fabian Isensee, Jens Petersen, Andre Klein, David Zimmerer, Paul F. Jaeger, Simon Kohl, Jakob Wasserthal, Gregor Koehler, Tobias Norajitra, Sebastian Wirkert, Klaus H. Maier-Hein
TL;DR
Medical image segmentation benchmarks often require interdependent, dataset-specific choices, complicating claims of general superiority. nnU-Net builds a self-adapting framework around simple 2D and 3D U-Nets, automatically configuring the surrounding pipeline; it achieves the challenge’s highest mean Dice scores across all classes and seven phase 1 tasks except BrainTumour class 1.
Problem
Dataset-specific architecture and training choices, together with limited validation across datasets, make generalizable segmentation performance difficult to assess.
Method
nnU-Net uses minimally modified 2D and 3D U-Net models with automatically adapted architectures, preprocessing, training, inference, and post-processing.
Results
nnU-Net achieves the highest mean Dice scores across all classes and seven phase 1 tasks except BrainTumour class 1 in the Medical Segmentation Decathlon leaderboard.
Takeaways & Limitations
The results support focusing on automatic, non-architectural design choices around simple U-Nets when adapting segmentation methods across datasets.
Takeaways & Limitations
The added benefit of several design choices was not properly validated in the challenge, motivating systematic ablation studies.
Abstract
from arXiv · showhide
The U-Net was presented in 2015. With its straight-forward and successful architecture it quickly evolved to a commonly used benchmark in medical image segmentation. The adaptation of the U-Net to novel problems, however, comprises several degrees of freedom regarding the exact architecture, preprocessing, training and inference. These choices are not independent of each other and substantially impact the overall performance. The present paper introduces the nnU-Net ('no-new-Net'), which refers to a robust and self-adapting framework on the basis of 2D and 3D vanilla U-Nets. We argue the strong case for taking away superfluous bells and whistles of many proposed network designs and instead focus on the remaining aspects that make out the performance and generalizability of a method. We evaluate the nnU-Net in the context of the Medical Segmentation Decathlon challenge, which measures segmentation performance in ten disciplines comprising distinct entities, image modalities, image geometries and dataset sizes, with no manual adjustments between datasets allowed. At the time of manuscript submission, nnU-Net achieves the highest mean dice scores across all classes and seven phase 1 tasks (except class 1 in BrainTumour) in the online leaderboard of the challenge.
1 Introduction
Medical image segmentation methods often require dataset-specific architecture and training changes, making broad claims difficult to validate. nnU-Net addresses this by automatically adapting U-Net-based architectures and the surrounding pipeline across diverse datasets.
- Motivation: Specialized architectures and training modifications are commonly required to achieve competitive performance across segmentation benchmarks.This proliferation of adaptations is accompanied by limited validation on only a few datasets.
- Motivation: The Medical Segmentation Decathlon evaluates whether one automatically adapting algorithm generalizes across ten datasets and three undisclosed phase 2 tasks.Participants may adapt to each dataset only fully automatically, and the final method is frozen before held-out evaluation.
- Motivation: Interdependent choices in architecture, preprocessing, training, inference, and post-processing can cause a standard U-Net benchmark to underperform on new problems.The paper hypothesizes that some architectural modifications are overfit to specific problems or evaluated using imperfect reimplementations.
- Contribution: nnU-Net uses minor U-Net modifications while automatically adapting image-geometry-dependent architectures and defining preprocessing, training, inference, and post-processing.The framework includes resampling and normalization, loss and optimization settings, augmentation, patch-based inference, ensembling, and optional connected-component post-processing.
2 Methods
The framework uses a small pool of basic U-Net architectures and focuses on automatically designed training and processing choices. Its encoder-decoder structure preserves spatial detail through skip connections, with limited architectural changes.
- Architecture: The framework keeps architectural modifications close to negligible and instead emphasizes an automatic training pipeline.This design shifts attention from adding architectural components to configuring the surrounding procedure.
- U-Net architecture: U-Net combines an encoder that aggregates semantic information with a decoder that restores spatial information using skip connections.The decoder recombines bottom-level semantic features with higher-resolution encoder maps to support fine-structure segmentation.
- Implementation: The implementation replaces ReLU with leaky ReLU and batch normalization with instance normalization while retaining plain convolutional and transposed-convolution layers.The leaky ReLU negative slope is 1e-2.
2D U-Net
The paper notes a trade-off in applying 2D or conventional 3D segmentation to anisotropic medical images. The 2D choice loses z-axis aggregation, while 3D methods can deteriorate on anisotropic data.
- 2D U-Net: A 2D U-Net cannot aggregate valuable information along the z-axis in 3D medical image segmentation.This makes the 2D formulation appear suboptimal when volumetric context is important.
- 2D U-Net: Conventional 3D segmentation methods can deteriorate when the dataset is anisotropic.The paper cites the Prostate dataset as an example of this setting.
- 2D U-Net: The comparison motivates retaining both 2D and 3D model options rather than assuming one dimensionality is optimal for every dataset.This is a direct implication of the two limitations described in the passage.
3D U-Net
The 3D U-Net is effective for volumetric data but patch-based training can limit context on large images. nnU-Net addresses this with a cascaded low-resolution-to-full-resolution model and automatic, geometry- and memory-aware configuration.
- 3D U-Net: Patch-based 3D U-Net training can limit the field of view and contextual information available for large images.GPU memory prevents training on entire patient images, and this can hinder distinguishing parts of a liver from other organs.
- 3D U-Net: The U-Net Cascade first segments downsampled images, then supplies upsampled one-hot segmentations as additional channels to a full-resolution 3D U-Net.The second stage refines the first-stage output on full-resolution patches.
- Dynamic adaptation of network topologies: Input patch size and pooling operations are automatically adapted to each dataset’s median shape and image geometry, subject to hardware constraints.The paper contrasts median shapes of 482 × 512 × 512 for Liver and 36 × 50 × 35 for Hippocampus.
- Dynamic adaptation of network topologies: The optimizer-step voxel budget is capped at 5% of the dataset, with batch size reduced when the limit is exceeded.The batch-size lower bound is 2.
- Dynamic adaptation of network topologies: The phase 1 topology table reports automatically generated configurations for seven tasks, with 3D U-Net lowres denoting the cascade’s first stage.The cascade’s second-stage configuration is identical to the 3D U-Net configuration.
2.2 Preprocessing
nnU-Net makes preprocessing part of a fully automated segmentation pipeline. It crops nonzero regions and resamples patients to address heterogeneous voxel spacings.
- Preprocessing is performed automatically without user intervention.
- Cropping to nonzero regions can reduce computational burden, especially for skull-stripped brain MRI.
- Patients are resampled to the median voxel spacing so networks can learn spatial semantics despite heterogeneous spacings.
Resampling
The framework combines dataset-dependent resampling, normalization, loss computation, optimization, and augmentation choices within its training procedure.
- Image data uses third-order spline interpolation, while segmentation masks use nearest-neighbor interpolation during resampling.
- The U-Net Cascade is triggered when median resampled shape exceeds four times the voxels processable by a 3D U-Net with batch size 2.
- CT images are clipped to the [0.5, 99.5] intensity percentiles and then z-score normalized using training-dataset statistics.
- Training combines Dice and crossentropy loss, with Dice computation differing between near-entire-patient networks and other networks.
- The Dice loss is a multi-class adaptation of a previously proposed formulation, using softmax outputs and one-hot ground-truth encodings.
- Adam is used with an initial learning rate of 3 × 10^-4 for all experiments.
- Training uses random rotations, scaling, elastic deformations, gamma correction, and mirroring to address overfitting with limited data.
Data Augmentation
nnU-Net uses fixed augmentation parameter sets separately for 2D and 3D U-Nets, with adaptations for geometric conditions and cascade inputs.
- Augmentation parameters are defined separately for 2D and 3D U-Nets and are not modified between datasets.
- When 3D patch geometry is highly unequal, the framework applies 2D augmentation slice-wise instead.
- Cascade segmentations receive random erosion, dilation, opening, closing, and connected-component removal to prevent strong co-adaptation.
- More than one-third of each batch must contain at least one randomly chosen foreground class to stabilize training.
2.4 Inference
Inference aggregates overlapping patch predictions with center weighting and mirroring-based test-time augmentation, then uses postprocessing and validation-selected model combinations.
- Overlapping patches are aggregated with higher weights for voxels near patch centers, and all valid-axis mirrors provide test-time augmentation.
- Tiled prediction and test-time augmentation can yield up to 64 predictions for a voxel, while five cross-validation networks are ensembled.
- Connected-component analysis removes all but the largest predicted component when a class is single-component in every training case.
- The final model or ensemble is selected by the highest mean foreground Dice score on training-set cross-validation.
3 Experiments and Results
The phase 1 evaluation used five-fold cross-validation and a single held-out test submission. Cross-validation performance was generally recovered on held-out data, suggesting limited over-fitting except for BrainTumour.
- Evaluation design: Five-fold cross-validation was used to optimize network topologies on the phase 1 datasets.The final submission model was highlighted, and the authors report making only one test-set submission.
- Held-out performance: Phase 1 cross-validation results were robustly recovered on held-out test sets, indicating an absence of over-fitting.BrainTumour was the exception, showing a performance dip across all foreground classes.
- Held-out performance: BrainTumour showed performance drops between validation and testing across all foreground classes.The authors attribute these drops to shifts in the data and/or ground-truth distributions, consistent with patterns reported for the BRATS challenge.
4 Discussion
nnU-Net builds on the original U-Net while automatically adapting the segmentation pipeline to each dataset. In the Medical Segmentation Decathlon, it achieved leading mean Dice scores across the evaluated phase 1 tasks, although several design choices lacked challenge-specific validation.
- Framework: nnU-Net dynamically adapts the original U-Net framework to each dataset.Its central contribution is automatic design of preprocessing, training, and inference rather than extensive architectural modification.
- Framework: The approach focuses on an automatic training pipeline while keeping architectural modifications to the original U-Net close to negligible.The framework uses comparatively simple U-Net models and omits several proposed architectural extensions.
- Results: The framework achieved the highest mean Dice scores for all classes in seven phase 1 tasks except class 1 in BrainTumour.This outcome was reported for the Medical Segmentation Decathlon online leaderboard at manuscript submission.
- Limitations: The added benefits of several design choices, including Leaky ReLUs and data-augmentation parameters, were not properly validated in the challenge.The authors identify systematic ablation studies as future work.