To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the nnUNet model on the Medical Segmentation Decathlon dataset. For the specifics on training and inference, see the Advanced section.

1. Clone the repository.

Executing this command will create your local repository with all the code to run nnU-Net.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/Segmentation/nnUNet

2. Build the nnU-Net PyTorch NGC container.

This command will use the Dockerfile to create a Docker image named nnunet, downloading all the required components automatically.

docker build -t nnunet .

The NGC container contains all the components optimized for usage on NVIDIA hardware.

3. Start an interactive session in the NGC container to run preprocessing/training/inference.

The following command will launch the container and mount the ./data directory as a volume to the /data directory inside the container, and ./results directory to the /results directory in the container.

mkdir data results
docker run -it --privileged --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/data:/data -v ${PWD}/results:/results nnunet:latest /bin/bash

4. Prepare the BraTS dataset.

To download and preprocess the data run:

python download.py --task 01
python preprocess.py --task 01 --dim 3
python preprocess.py --task 01 --dim 2

Then ls /data should print:

01_3d 01_2d Task01_BrainTumour

For the specifics on data preprocessing, see the Getting the data section.

5. Start training.

Training can be started with:

python scripts/train.py --gpus <gpus> --fold <fold> --dim <dim> [--amp] [--bind]

To see descriptions of the train script arguments run python scripts/train.py --help. You can customize the training process. For details, see the Training process section.

6. Start benchmarking.

The training and inference performance can be evaluated by using benchmarking scripts, such as:

python scripts/benchmark.py --mode {train,predict} --gpus <ngpus> --dim {2,3} --batch_size <bsize> [--amp] [--bind]

To see descriptions of the benchmark script arguments run python scripts/benchmark.py --help.

7. Start inference/predictions.

Inference can be started with:

python scripts/inference.py --data <path/to/data> --dim <dim> --fold <fold> --ckpt_path <path/to/checkpoint> [--amp] [--tta] [--save_preds]

Note: You have to prepare either validation or test dataset to run this script by running python preprocess.py --task 01 --dim {2,3} --exec_mode {val,test}. After preprocessing inside given task directory (e.g. /data/01_3d/ for task 01 and dim 3) it will create val or test directory with preprocessed data ready for inference. Possible workflow:

python preprocess.py --task 01 --dim 3 --exec_mode val
python scripts/inference.py --data /data/01_3d/val --dim 3 --fold 0 --ckpt_path <path/to/checkpoint> --amp --tta --save_preds

Then if you have labels for predicted images you can evaluate them with evaluate.py script. For example:

python evaluate.py --preds /results/preds_task_01_dim_3_fold_0_tta --lbls /data/Task01_BrainTumour/labelsTr

To see descriptions of the inference script arguments run python scripts/inference.py --help. You can customize the inference process. For details, see the Inference process section.

Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark yours performance to Training performance benchmark, or Inference performance benchmark. Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the Results section.