The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

• main.py: Entry point to the application. Runs training, evaluation, inference or benchmarking.
• preprocess.py: Entry point to data preprocessing.
• download.py: Downloads given dataset from Medical Segmentation Decathlon.
• Dockerfile: Container with the basic set of dependencies to run nnU-Net.
• requirements.txt: Set of extra requirements for running nnU-Net.
• evaluate.py: Compare predictions with ground truth and get the final score.

The data_preprocessing folder contains information about the data preprocessing used by nnU-Net. Its contents are:

• configs.py: Defines dataset configuration like patch size or spacing.
• preprocessor.py: Implements data preprocessing pipeline.

The data_loading folder contains information about the data-loading pipeline used by nnU-Net. Its contents are:

• data_module.py: Defines a data module managing datasets and splits (similar to PyTorch Lightning DataModule)
• dali_loader.py: Implements DALI data loading pipelines.
• utils.py: Defines auxiliary functions used for data loading.

The models folder contains information about the building blocks of nnU-Net and the way they are assembled. Its contents are:

• layers.py: Implements convolution blocks used by the U-Net template.
• nn_unet.py: Implements training/validation/test logic and dynamic creation of U-Net architecture used by nnU-Net.
• sliding_window.py: Implements sliding window inference used by evaluation and prediction loops.
• unet.py: Implements the U-Net template.

The runtime folder contains information about training, inference, and evaluation logic. Its contents are:

• args.py: Defines command line arguments.
• checkpoint.py: Implements checkpoint saving.
• logging.py: Defines logging utilities along with wandb.io integration.
• losses.py: Implements loss functions.
• metrics.py: Implements dice metric and metric management.
• run.py: Implements training loop and inference and evaluation logic.
• utils.py: Defines auxiliary functions used during runtime.

Other folders included in the root directory are:

• images/: Contains a model diagram.
• scripts/: Provides scripts for training, benchmarking, and inference of nnU-Net.

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command-line option, for example:

python main.py --help

The following example output is printed when running the model:

usage: main.py [-h] [--exec-mode {train,evaluate,predict,export,nav}] [--gpus GPUS] [--data DATA] [--task TASK] [--dim {2,3}]
               [--seed SEED] [--benchmark] [--tta [BOOLEAN]] [--save-preds [BOOLEAN]] [--sw-benchmark [BOOLEAN]]
               [--results RESULTS] [--logname LOGNAME] [--quiet] [--use-dllogger [BOOLEAN]] [--amp [BOOLEAN]] [--xla [BOOLEAN]]
               [--read-roi [BOOLEAN]] [--batch-size BATCH_SIZE] [--learning-rate LEARNING_RATE] [--momentum MOMENTUM]
               [--scheduler {none,poly,cosine,cosine_annealing}] [--end-learning-rate END_LEARNING_RATE]
               [--cosine-annealing-first-cycle-steps COSINE_ANNEALING_FIRST_CYCLE_STEPS]
               [--cosine-annealing-peak-decay COSINE_ANNEALING_PEAK_DECAY] [--optimizer {sgd,adam,radam}]
               [--deep-supervision [BOOLEAN]] [--lookahead [BOOLEAN]] [--weight-decay WEIGHT_DECAY]
               [--loss-batch-reduction [BOOLEAN]] [--loss-include-background [BOOLEAN]] [--negative-slope NEGATIVE_SLOPE]
               [--norm {instance,batch,group,none}] [--ckpt-strategy {last_and_best,last_only,none}] [--ckpt-dir CKPT_DIR]
               [--saved-model-dir SAVED_MODEL_DIR] [--resume-training] [--load_sm [BOOLEAN]] [--validate [BOOLEAN]] [--nvol NVOL]
               [--oversampling OVERSAMPLING] [--num-workers NUM_WORKERS] [--sw-batch-size SW_BATCH_SIZE] [--overlap OVERLAP]
               [--blend {gaussian,constant}] [--nfolds NFOLDS] [--fold FOLD] [--epochs EPOCHS] [--skip-eval SKIP_EVAL]
               [--steps-per-epoch STEPS_PER_EPOCH] [--bench-steps BENCH_STEPS] [--warmup-steps WARMUP_STEPS]

optional arguments:
  -h, --help            show this help message and exit
  --exec-mode {train,evaluate,predict,export,nav}, --exec_mode {train,evaluate,predict,export,nav}
                        Execution mode to run the model (default: train)
  --gpus GPUS
  --data DATA           Path to data directory (default: /data)
  --task TASK           Task number, MSD uses numbers 01-10 (default: 01)
  --dim {2,3}           UNet dimension (default: 3)
  --seed SEED           Random seed (default: None)
  --benchmark           Run model benchmarking (default: False)
  --tta [BOOLEAN]       Enable test time augmentation (default: False)
  --save-preds [BOOLEAN], --save_preds [BOOLEAN]
                        Save predictions (default: False)
  --sw-benchmark [BOOLEAN], --sw_benchmark [BOOLEAN]
  --results RESULTS     Path to results directory (default: /results)
  --logname LOGNAME     DLLogger output filename (default: dllogger.json)
  --quiet               Minimalize stdout/stderr output (default: False)
  --use-dllogger [BOOLEAN], --use_dllogger [BOOLEAN]
                        Use DLLogger logging (default: True)
  --amp [BOOLEAN]       Enable automatic mixed precision (default: False)
  --xla [BOOLEAN]       Enable XLA compiling (default: False)
  --read-roi [BOOLEAN], --read_roi [BOOLEAN]
                        Use DALI direct ROI loading feature (default: False)
  --batch-size BATCH_SIZE, --batch_size BATCH_SIZE
                        Batch size (default: 2)
  --learning-rate LEARNING_RATE, --learning_rate LEARNING_RATE
                        Learning rate (default: 0.0003)
  --momentum MOMENTUM   Momentum factor (SGD only) (default: 0.99)
  --scheduler {none,poly,cosine,cosine_annealing}
                        Learning rate scheduler (default: none)
  --end-learning-rate END_LEARNING_RATE
                        End learning rate for poly scheduler (default: 5e-05)
  --cosine-annealing-first-cycle-steps COSINE_ANNEALING_FIRST_CYCLE_STEPS
                        Length of a cosine decay cycle in steps, only with 'cosine_annealing' scheduler (default: 512)
  --cosine-annealing-peak-decay COSINE_ANNEALING_PEAK_DECAY
                        Multiplier reducing initial learning rate (default: 0.95)
  --optimizer {sgd,adam,radam}
                        Optimizer (default: adam)
  --deep-supervision [BOOLEAN], --deep_supervision [BOOLEAN]
                        Use deep supervision. (default: False)
  --lookahead [BOOLEAN]
                        Use Lookahead with the optimizer (default: False)
  --weight-decay WEIGHT_DECAY, --weight_decay WEIGHT_DECAY
                        Weight decay (L2 penalty) (default: 0.0001)
  --loss-batch-reduction [BOOLEAN]
                        Reduce batch dimension first during loss calculation (default: True)
  --loss-include-background [BOOLEAN]
                        Include background class to loss calculation (default: False)
  --negative-slope NEGATIVE_SLOPE
                        Negative slope for LeakyReLU (default: 0.01)
  --norm {instance,batch,group,none}
                        Type of normalization layers (default: instance)
  --ckpt-strategy {last_and_best,last_only,none}
                        Strategy how to save checkpoints (default: last_and_best)
  --ckpt-dir CKPT_DIR   Path to checkpoint directory (default: /results/ckpt)
  --saved-model-dir SAVED_MODEL_DIR
                        Path to saved model directory (for evaluation and prediction) (default: None)
  --resume-training, --resume_training
                        Resume training from the last checkpoint (default: False)
  --load_sm [BOOLEAN]   Load exported savedmodel (default: False)
  --validate [BOOLEAN]  Validate exported savedmodel (default: False)
  --nvol NVOL           Number of volumes which come into single batch size for 2D model (default: 2)
  --oversampling OVERSAMPLING
                        Probability of crop to have some region with positive label (default: 0.33)
  --num-workers NUM_WORKERS
                        Number of subprocesses to use for data loading (default: 8)
  --sw-batch-size SW_BATCH_SIZE
                        Sliding window inference batch size (default: 2)
  --overlap OVERLAP     Amount of overlap between scans during sliding window inference (default: 0.5)
  --blend {gaussian,constant}, --blend-mode {gaussian,constant}
                        How to blend output of overlapping windows (default: gaussian)
  --nfolds NFOLDS       Number of cross-validation folds (default: 5)
  --fold FOLD           Fold number (default: 0)
  --epochs EPOCHS       Number of epochs (default: 1000)
  --skip-eval SKIP_EVAL
                        Skip evaluation for the first N epochs. (default: 0)
  --steps-per-epoch STEPS_PER_EPOCH
                        Steps per epoch. By default ceil(training_dataset_size / batch_size / gpus) (default: None)
  --bench-steps BENCH_STEPS
                        Number of benchmarked steps in total (default: 100)
  --warmup-steps WARMUP_STEPS
                        Number of warmup steps before collecting benchmarking statistics (default: 25)

Getting the data

The nnU-Net model was trained on the Medical Segmentation Decathlon datasets. All datasets are in Neuroimaging Informatics Technology Initiative (NIfTI) format.

Dataset guidelines

To train nnU-Net you will need to preprocess your dataset as the first step with preprocess.py script. Run python scripts/preprocess.py --help to see descriptions of the preprocess script arguments.

For example to preprocess data for 3D U-Net run: python preprocess.py --task 01 --dim 3.

In data_preprocessing/configs.py for each Medical Segmentation Decathlon task, there are defined: patch sizes, precomputed spacings and statistics for CT datasets.

The preprocessing pipeline consists of the following steps:

1. Cropping to the region of non-zero values.
2. Resampling to the median voxel spacing of their respective dataset (exception for anisotropic datasets where the lowest resolution axis is selected to be the 10th percentile of the spacings).
3. Padding volumes so that dimensions are at least as patch size.
4. Normalizing:

• For CT modalities the voxel values are clipped to 0.5 and 99.5 percentiles of the foreground voxels and then data is normalized with mean and standard deviation collected from foreground voxels.
• For MRI modalities z-score normalization is applied.

Multi-dataset

It is possible to run nnUNet on a custom dataset. If your dataset corresponds to Medical Segmentation Decathlon (i.e. data should be in NIfTi format and there should be dataset.json file where you need to provide fields: modality, labels, and at least one of training, test) you need to perform the following:

1. Mount your dataset to /data directory.

2. In data_preprocessing/config.py:

• Add to the task_dir dictionary your dataset directory name. For example, for the Brain Tumour dataset, it corresponds to "01": "Task01_BrainTumour".
• Add the patch size that you want to use for training to the patch_size dictionary. For example, for the Brain Tumour dataset it corresponds to "01_3d": [128, 128, 128] for 3D U-Net and "01_2d": [192, 160] for 2D U-Net. There are three types of suffixes _3d, _2d corresponding to 3D UNet and 2D U-Net.

3. Preprocess your data with preprocess.py scripts. For example, to preprocess the Brain Tumour dataset for 2D U-Net you should run python preprocess.py --task 01 --dim 2.

Training process

The model trains for at most --epochs epochs. After each epoch evaluation, the validation set is done and validation metrics are monitored for checkpoint updating (see --ckpt-strategy flag). Default training settings are:

• The Adam optimizer with a learning rate of 0.0003 and weight decay of 0.0001.
• Training batch size is set to 2.

This default parametrization is applied when running scripts from the scripts/ directory and when running main.py without explicitly overriding these parameters. By default, using scripts from the scripts directory will not use AMP. To enable AMP, pass the --amp flag. AMP can be enabled for every mode of execution. However, a custom invocation of the main.py script will turn on AMP by default. To turn it off use main.py --amp false.

The default configuration minimizes a function L = (1 - dice_coefficient) + cross_entropy during training and reports achieved convergence as dice coefficient per class. The training, with a combination of dice and cross-entropy has been proven to achieve better convergence than training using only dice.

The training can be run without using the predefined scripts. The name of the training script is main.py. For example:

python main.py --exec-mode train --task 01 --fold 0

Training artifacts will be saved to /results in the container. Some important artifacts are:

• /results/dllogger.json: Collected dice scores and loss values evaluated after each epoch during training on a validation set.
• /results/ckpt/: Saved checkpoints. By default, two checkpoints are saved - one after each epoch and one with the highest validation dice (saved in the /results/ckpt/best/ subdirectory). You can change this behavior by modifying the --ckpt-strategy parameter.

To load the pretrained model, provide --ckpt-dir <path/to/checkpoint/directory> and use --resume-training if you intend to continue training.

To use multi-GPU training with the main.py script prepend the command with horovodrun -np <ngpus>, for example with 8 GPUs use:

horovodrun -np 8 python main.py --exec-mode train --task 01 --fold 0

Inference process

Inference can be launched by passing the --exec-mode predict flag. For example:

python main.py --exec-mode predict --xla --task 01 --fold 0 --gpus 1 --amp --tta --save-preds --ckpt-dir <path/to/checkpoint/dir>

The script will then:

• Load the checkpoint from the directory specified by the <path/to/checkpoint/dir> directory
• Run inference on the preprocessed validation dataset corresponding to fold 0
• If --save-preds is provided then resulting masks in the NumPy format will be saved in the /results directory