nnU-Net requires some environment variables so that it always knows where the raw data, preprocessed data and trained
models are. Depending on the operating system, these environment variables need to be set in different ways.
Variables can either be set permanently (recommended!) or you can decide to set them everytime you call nnU-Net.
Locate the .bashrc
file in your home folder and add the following lines to the bottom:
export nnUNet_raw="/media/fabian/nnUNet_raw"
export nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed"
export nnUNet_results="/media/fabian/nnUNet_results"
(of course you need to adapt the paths to the actual folders you intend to use).
If you are using a different shell, such as zsh, you will need to find the correct script for it. For zsh this is .zshrc
.
Just execute the following lines whenever you run nnU-Net:
export nnUNet_raw="/media/fabian/nnUNet_raw"
export nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed"
export nnUNet_results="/media/fabian/nnUNet_results"
(of course you need to adapt the paths to the actual folders you intend to use).
Important: These variables will be deleted if you close your terminal! They will also only apply to the current
terminal window and DO NOT transfer to other terminals!
Alternatively you can also just prefix them to your nnU-Net commands:
nnUNet_results="/media/fabian/nnUNet_results" nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed" nnUNetv2_train[...]
You can always execute echo ${nnUNet_raw}
etc to print the environment variables. This will return an empty string if
they were not set.
Useful links:
See Set Environment Variable in Windows via GUI
here.
Or read about setx (command prompt).
Just execute the following before you run nnU-Net:
(powershell)
$Env:nnUNet_raw = "/media/fabian/nnUNet_raw"
$Env:nnUNet_preprocessed = "/media/fabian/nnUNet_preprocessed"
$Env:nnUNet_results = "/media/fabian/nnUNet_results"
(command prompt)
set nnUNet_raw="/media/fabian/nnUNet_raw"
set nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed"
set nnUNet_results="/media/fabian/nnUNet_results"
(of course you need to adapt the paths to the actual folders you intend to use).
Important: These variables will be deleted if you close your session! They will also only apply to the current
window and DO NOT transfer to other sessions!
Printing in Windows works differently depending on the environment you are in:
powershell: echo $Env:[variable_name]
command prompt: echo %[variable_name]%
The only way to bring your data into nnU-Net is by storing it in a specific format. Due to nnU-Net’s roots in the
Medical Segmentation Decathlon (MSD), its dataset is heavily inspired but has since
diverged (see also here) from the format used in the MSD.
Datasets consist of three components: raw images, corresponding segmentation maps and a dataset.json file specifying
some metadata.
If you are migrating from nnU-Net v1, read this to convert your existing Tasks.
Each training case is associated with an identifier = a unique name for that case. This identifier is used by nnU-Net to
connect images with the correct segmentation.
A training case consists of images and their corresponding segmentation.
Images is plural because nnU-Net supports arbitrarily many input channels. In order to be as flexible as possible,
nnU-net requires each input channel to be stored in a separate image (with the sole exception being RGB natural
images). So these images could for example be a T1 and a T2 MRI (or whatever else you want). The different input
channels MUST have the same geometry (same shape, spacing (if applicable) etc.) and
must be co-registered (if applicable). Input channels are identified by nnU-Net by their FILE_ENDING: a four-digit integer at the end
of the filename. Image files must therefore follow the following naming convention: {CASE_IDENTIFIER}_{XXXX}.{FILE_ENDING}.
Hereby, XXXX is the 4-digit modality/channel identifier (should be unique for each modality/chanel, e.g., “0000” for T1, “0001” for
T2 MRI, …) and FILE_ENDING is the file extension used by your image format (.png, .nii.gz, …). See below for concrete examples.
The dataset.json file connects channel names with the channel identifiers in the ‘channel_names’ key (see below for details).
Side note: Typically, each channel/modality needs to be stored in a separate file and is accessed with the XXXX channel identifier.
Exception are natural images (RGB; .png) where the three color channels can all be stored in one file (see the road segmentation dataset as an example).
Segmentations must share the same geometry with their corresponding images (same shape etc.). Segmentations are
integer maps with each value representing a semantic class. The background must be 0. If there is no background, then
do not use the label 0 for something else! Integer values of your semantic classes must be consecutive (0, 1, 2, 3,
…). Of course, not all labels have to be present in each training case. Segmentations are saved as {CASE_IDENTIFER}.{FILE_ENDING} .
Within a training case, all image geometries (input channels, corresponding segmentation) must match. Between training
cases, they can of course differ. nnU-Net takes care of that.
Important: The input channels must be consistent! Concretely, all images need the same input channels in the same
order and all input channels have to be present every time. This is also true for inference!
nnU-Net expects the same file format for images and segmentations! These will also be used for inference. For now, it
is thus not possible to train .png and then run inference on .jpg.
One big change in nnU-Net V2 is the support of multiple input file types. Gone are the days of converting everything to .nii.gz!
This is implemented by abstracting the input and output of images + segmentations through BaseReaderWriter
. nnU-Net
comes with a broad collection of Readers+Writers and you can even add your own to support your data format!
See here.
As a nice bonus, nnU-Net now also natively supports 2D input images and you no longer have to mess around with
conversions to pseudo 3D niftis. Yuck. That was disgusting.
Note that internally (for storing and accessing preprocessed images) nnU-Net will use its own file format, irrespective
of what the raw data was provided in! This is for performance reasons.
By default, the following file formats are supported:
The file extension lists are not exhaustive and depend on what the backend supports. For example, nibabel and SimpleITK
support more than the three given here. The file endings given here are just the ones we tested!
IMPORTANT: nnU-Net can only be used with file formats that use lossless (or no) compression! Because the file
format is defined for an entire dataset (and not separately for images and segmentations, this could be a todo for
the future), we must ensure that there are no compression artifacts that destroy the segmentation maps. So no .jpg and
the likes!
Datasets must be located in the nnUNet_raw
folder (which you either define when installing nnU-Net or export/set every
time you intend to run nnU-Net commands!).
Each segmentation dataset is stored as a separate ‘Dataset’. Datasets are associated with a dataset ID, a three digit
integer, and a dataset name (which you can freely choose): For example, Dataset005_Prostate has ‘Prostate’ as dataset name and
the dataset id is 5. Datasets are stored in the nnUNet_raw
folder like this:
nnUNet_raw/
├── Dataset001_BrainTumour
├── Dataset002_Heart
├── Dataset003_Liver
├── Dataset004_Hippocampus
├── Dataset005_Prostate
├── ...
Within each dataset folder, the following structure is expected:
Dataset001_BrainTumour/
├── dataset.json
├── imagesTr
├── imagesTs # optional
└── labelsTr
When adding your custom dataset, take a look at the dataset_conversion folder and
pick an id that is not already taken. IDs 001-010 are for the Medical Segmentation Decathlon.
The scheme introduced above results in the following folder structure. Given
is an example for the first Dataset of the MSD: BrainTumour. This dataset hat four input channels: FLAIR (0000),
T1w (0001), T1gd (0002) and T2w (0003). Note that the imagesTs folder is optional and does not have to be present.
nnUNet_raw/Dataset001_BrainTumour/
├── dataset.json
├── imagesTr
│ ├── BRATS_001_0000.nii.gz
│ ├── BRATS_001_0001.nii.gz
│ ├── BRATS_001_0002.nii.gz
│ ├── BRATS_001_0003.nii.gz
│ ├── BRATS_002_0000.nii.gz
│ ├── BRATS_002_0001.nii.gz
│ ├── BRATS_002_0002.nii.gz
│ ├── BRATS_002_0003.nii.gz
│ ├── ...
├── imagesTs
│ ├── BRATS_485_0000.nii.gz
│ ├── BRATS_485_0001.nii.gz
│ ├── BRATS_485_0002.nii.gz
│ ├── BRATS_485_0003.nii.gz
│ ├── BRATS_486_0000.nii.gz
│ ├── BRATS_486_0001.nii.gz
│ ├── BRATS_486_0002.nii.gz
│ ├── BRATS_486_0003.nii.gz
│ ├── ...
└── labelsTr
├── BRATS_001.nii.gz
├── BRATS_002.nii.gz
├── ...
Here is another example of the second dataset of the MSD, which has only one input channel:
nnUNet_raw/Dataset002_Heart/
├── dataset.json
├── imagesTr
│ ├── la_003_0000.nii.gz
│ ├── la_004_0000.nii.gz
│ ├── ...
├── imagesTs
│ ├── la_001_0000.nii.gz
│ ├── la_002_0000.nii.gz
│ ├── ...
└── labelsTr
├── la_003.nii.gz
├── la_004.nii.gz
├── ...
Remember: For each training case, all images must have the same geometry to ensure that their pixel arrays are aligned. Also
make sure that all your data is co-registered!
See also dataset format inference!!
The dataset.json contains metadata that nnU-Net needs for training. We have greatly reduced the number of required
fields since version 1!
Here is what the dataset.json should look like at the example of the Dataset005_Prostate from the MSD:
{
"channel_names": { # formerly modalities
"0": "T2",
"1": "ADC"
},
"labels": { # THIS IS DIFFERENT NOW!
"background": 0,
"PZ": 1,
"TZ": 2
},
"numTraining": 32,
"file_ending": ".nii.gz"
"overwrite_image_reader_writer": "SimpleITKIO" # optional! If not provided nnU-Net will automatically determine the ReaderWriter
}
The channel_names determine the normalization used by nnU-Net. If a channel is marked as ‘CT’, then a global
normalization based on the intensities in the foreground pixels will be used. If it is something else, per-channel
z-scoring will be used. Refer to the methods section in our paper
for more details. nnU-Net v2 introduces a few more normalization schemes to
choose from and allows you to define your own, see here for more information.
Important changes relative to nnU-Net v1:
There is a utility with which you can generate the dataset.json automatically. You can find it
here.
See our examples in dataset_conversion for how to use it. And read its documentation!
If you are migrating from the old nnU-Net, convert your existing datasets with nnUNetv2_convert_old_nnUNet_dataset
!
Example for migrating a nnU-Net v1 Task:
nnUNetv2_convert_old_nnUNet_dataset /media/isensee/raw_data/nnUNet_raw_data_base/nnUNet_raw_data/Task027_ACDC Dataset027_ACDC
Use nnUNetv2_convert_old_nnUNet_dataset -h
for detailed usage instructions.
See convert_msd_dataset.md
2D is now natively supported (yay!). See here as well as the example dataset in this
script.
When updating a dataset it is best practice to remove the preprocessed data in nnUNet_preprocessed/DatasetXXX_NAME
to ensure a fresh start. Then replace the data in nnUNet_raw
and rerun nnUNetv2_plan_and_preprocess
. Optionally,
also remove the results from old trainings.
In the dataset_conversion
folder (see here) are multiple example scripts for
converting datasets into nnU-Net format. These scripts cannot be run as they are (you need to open them and change
some paths) but they are excellent examples for you to learn how to convert your own datasets into nnU-Net format.
Just pick the dataset that is closest to yours as a starting point.
The list of dataset conversion scripts is continually updated. If you find that some publicly available dataset is
missing, feel free to open a PR to add it!
nnU-Net relies on environment variables to know where raw data, preprocessed data and trained model weights are stored.
To use the full functionality of nnU-Net, the following three environment variables must be set:
nnUNet_raw
: This is where you place the raw datasets. This folder will have one subfolder for each dataset names
DatasetXXX_YYY where XXX is a 3-digit identifier (such as 001, 002, 043, 999, …) and YYY is the (unique)
dataset name. The datasets must be in nnU-Net format, see here.
Example tree structure:
nnUNet_raw/Dataset001_NAME1
├── dataset.json
├── imagesTr
│ ├── ...
├── imagesTs
│ ├── ...
└── labelsTr
├── ...
nnUNet_raw/Dataset002_NAME2
├── dataset.json
├── imagesTr
│ ├── ...
├── imagesTs
│ ├── ...
└── labelsTr
├── ...
nnUNet_preprocessed
: This is the folder where the preprocessed data will be saved. The data will also be read from
this folder during training. It is important that this folder is located on a drive with low access latency and high
throughput (such as a nvme SSD (PCIe gen 3 is sufficient)).
nnUNet_results
: This specifies where nnU-Net will save the model weights. If pretrained models are downloaded, this
is where it will save them.
Given some dataset, nnU-Net fully automatically configures an entire segmentation pipeline that matches its properties.
nnU-Net covers the entire pipeline, from preprocessing to model configuration, model training, postprocessing
all the way to ensembling. After running nnU-Net, the trained model(s) can be applied to the test cases for inference.
nnU-Net expects datasets in a structured format. This format is inspired by the data structure of
the Medical Segmentation Decthlon. Please read
this for information on how to set up datasets to be compatible with nnU-Net.
Since version 2 we support multiple image file formats (.nii.gz, .png, .tif, …)! Read the dataset_format
documentation to learn more!
Datasets from nnU-Net v1 can be converted to V2 by running nnUNetv2_convert_old_nnUNet_dataset INPUT_FOLDER OUTPUT_DATASET_NAME
. Remember that v2 calls datasets DatasetXXX_Name (not Task) where XXX is a 3-digit number.
Please provide the path to the old task, not just the Task name. nnU-Net V2 doesn’t know where v1 tasks were!
Given a new dataset, nnU-Net will extract a dataset fingerprint (a set of dataset-specific properties such as
image sizes, voxel spacings, intensity information etc). This information is used to design three U-Net configurations.
Each of these pipelines operates on its own preprocessed version of the dataset.
The easiest way to run fingerprint extraction, experiment planning and preprocessing is to use:
nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity
Where DATASET_ID
is the dataset id (duh). We recommend --verify_dataset_integrity
whenever it’s the first time
you run this command. This will check for some of the most common error sources!
You can also process several datasets at once by giving -d 1 2 3 [...]
. If you already know what U-Net configuration
you need you can also specify that with -c 3d_fullres
(make sure to adapt -np in this case!). For more information
about all the options available to you please run nnUNetv2_plan_and_preprocess -h
.
nnUNetv2_plan_and_preprocess will create a new subfolder in your nnUNet_preprocessed folder named after the dataset.
Once the command is completed there will be a dataset_fingerprint.json file as well as a nnUNetPlans.json file for you to look at
(in case you are interested!). There will also be subfolders containing the preprocessed data for your UNet configurations.
[Optional]
If you prefer to keep things separate, you can also use nnUNetv2_extract_fingerprint
, nnUNetv2_plan_experiment
and nnUNetv2_preprocess
(in that order).
You pick which configurations (2d, 3d_fullres, 3d_lowres, 3d_cascade_fullres) should be trained! If you have no idea
what performs best on your data, just run all of them and let nnU-Net identify the best one. It’s up to you!
nnU-Net trains all configurations in a 5-fold cross-validation over the training cases. This is 1) needed so that
nnU-Net can estimate the performance of each configuration and tell you which one should be used for your
segmentation problem and 2) a natural way of obtaining a good model ensemble (average the output of these 5 models
for prediction) to boost performance.
You can influence the splits nnU-Net uses for 5-fold cross-validation (see here). If you
prefer to train a single model on all training cases, this is also possible (see below).
Note that not all U-Net configurations are created for all datasets. In datasets with small image sizes, the U-Net
cascade (and with it the 3d_lowres configuration) is omitted because the patch size of the full resolution U-Net
already covers a large part of the input images.
Training models is done with the nnUNetv2_train
command. The general structure of the command is:
nnUNetv2_train DATASET_NAME_OR_ID UNET_CONFIGURATION FOLD [additional options, see -h]
UNET_CONFIGURATION is a string that identifies the requested U-Net configuration (defaults: 2d, 3d_fullres, 3d_lowres,
3d_cascade_lowres). DATASET_NAME_OR_ID specifies what dataset should be trained on and FOLD specifies which fold of
the 5-fold-cross-validation is trained.
nnU-Net stores a checkpoint every 50 epochs. If you need to continue a previous training, just add a --c
to the
training command.
IMPORTANT: If you plan to use nnUNetv2_find_best_configuration
(see below) add the --npz
flag. This makes
nnU-Net save the softmax outputs during the final validation. They are needed for that. Exported softmax
predictions are very large and therefore can take up a lot of disk space, which is why this is not enabled by default.
If you ran initially without the --npz
flag but now require the softmax predictions, simply rerun the validation with:
nnUNetv2_train DATASET_NAME_OR_ID UNET_CONFIGURATION FOLD --val --npz
You can specify the device nnU-net should use by using -device DEVICE
. DEVICE can only be cpu, cuda or mps. If
you have multiple GPUs, please select the gpu id using CUDA_VISIBLE_DEVICES=X nnUNetv2_train [...]
(requires device to be cuda).
See nnUNetv2_train -h
for additional options.
For FOLD in [0, 1, 2, 3, 4], run:
nnUNetv2_train DATASET_NAME_OR_ID 2d FOLD [--npz]
For FOLD in [0, 1, 2, 3, 4], run:
nnUNetv2_train DATASET_NAME_OR_ID 3d_fullres FOLD [--npz]
For FOLD in [0, 1, 2, 3, 4], run:
nnUNetv2_train DATASET_NAME_OR_ID 3d_lowres FOLD [--npz]
For FOLD in [0, 1, 2, 3, 4], run:
nnUNetv2_train DATASET_NAME_OR_ID 3d_cascade_fullres FOLD [--npz]
Note that the 3D full resolution U-Net of the cascade requires the five folds of the low resolution U-Net to be
completed!
The trained models will be written to the nnUNet_results folder. Each training obtains an automatically generated
output folder name:
nnUNet_results/DatasetXXX_MYNAME/TRAINER_CLASS_NAME__PLANS_NAME__CONFIGURATION/FOLD
For Dataset002_Heart (from the MSD), for example, this looks like this:
nnUNet_results/
├── Dataset002_Heart
│── nnUNetTrainer__nnUNetPlans__2d
│ ├── fold_0
│ ├── fold_1
│ ├── fold_2
│ ├── fold_3
│ ├── fold_4
│ ├── dataset.json
│ ├── dataset_fingerprint.json
│ └── plans.json
└── nnUNetTrainer__nnUNetPlans__3d_fullres
├── fold_0
├── fold_1
├── fold_2
├── fold_3
├── fold_4
├── dataset.json
├── dataset_fingerprint.json
└── plans.json
Note that 3d_lowres and 3d_cascade_fullres do not exist here because this dataset did not trigger the cascade. In each
model training output folder (each of the fold_x folder), the following files will be created:
--npz
was set thenDuring training it is often useful to watch the progress. We therefore recommend that you have a look at the generated
progress.png when running the first training. It will be updated after each epoch.
Training times largely depend on the GPU. The smallest GPU we recommend for training is the Nvidia RTX 2080ti. With
that all network trainings take less than 2 days. Refer to our benchmarks to see if your system is
performing as expected.
If multiple GPUs are at your disposal, the best way of using them is to train multiple nnU-Net trainings at once, one
on each GPU. This is because data parallelism never scales perfectly linearly, especially not with small networks such
as the ones used by nnU-Net.
Example:
CUDA_VISIBLE_DEVICES=0 nnUNetv2_train DATASET_NAME_OR_ID 2d 0 [--npz] & # train on GPU 0
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train DATASET_NAME_OR_ID 2d 1 [--npz] & # train on GPU 1
CUDA_VISIBLE_DEVICES=2 nnUNetv2_train DATASET_NAME_OR_ID 2d 2 [--npz] & # train on GPU 2
CUDA_VISIBLE_DEVICES=3 nnUNetv2_train DATASET_NAME_OR_ID 2d 3 [--npz] & # train on GPU 3
CUDA_VISIBLE_DEVICES=4 nnUNetv2_train DATASET_NAME_OR_ID 2d 4 [--npz] & # train on GPU 4
...
wait
Important: The first time a training is run nnU-Net will extract the preprocessed data into uncompressed numpy
arrays for speed reasons! This operation must be completed before starting more than one training of the same
configuration! Wait with starting subsequent folds until the first training is using the GPU! Depending on the
dataset size and your System this should oly take a couple of minutes at most.
If you insist on running DDP multi-GPU training, we got you covered:
nnUNetv2_train DATASET_NAME_OR_ID 2d 0 [--npz] -num_gpus X
Again, note that this will be slower than running separate training on separate GPUs. DDP only makes sense if you have
manually interfered with the nnU-Net configuration and are training larger models with larger patch and/or batch sizes!
Important when using -num_gpus
:
In contrast to the old nnU-Net, DDP is now completely hassle free. Enjoy!
Once the desired configurations were trained (full cross-validation) you can tell nnU-Net to automatically identify
the best combination for you:
nnUNetv2_find_best_configuration DATASET_NAME_OR_ID -c CONFIGURATIONS
CONFIGURATIONS
hereby is the list of configurations you would like to explore. Per default, ensembling is enabled
meaning that nnU-Net will generate all possible combinations of ensembles (2 configurations per ensemble). This requires
the .npz files containing the predicted probabilities of the validation set to be present (use nnUNetv2_train
with
--npz
flag, see above). You can disable ensembling by setting the --disable_ensembling
flag.
See nnUNetv2_find_best_configuration -h
for more options.
nnUNetv2_find_best_configuration will also automatically determine the postprocessing that should be used.
Postprocessing in nnU-Net only considers the removal of all but the largest component in the prediction (once for
foreground vs background and once for each label/region).
Once completed, the command will print to your console exactly what commands you need to run to make predictions. It
will also create two files in the nnUNet_results/DATASET_NAME
folder for you to inspect:
inference_instructions.txt
again contains the exact commands you need to use for predictionsinference_information.json
can be inspected to see the performance of all configurations and ensembles, as wellRemember that the data located in the input folder must have the file endings as the dataset you trained the model on
and must adhere to the nnU-Net naming scheme for image files (see dataset format and
inference data format!)
nnUNetv2_find_best_configuration
(see above) will print a string to the terminal with the inference commands you need to use.
The easiest way to run inference is to simply use these commands.
If you wish to manually specify the configuration(s) used for inference, use the following commands:
For each of the desired configurations, run:
nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -c CONFIGURATION --save_probabilities
Only specify --save_probabilities
if you intend to use ensembling. --save_probabilities
will make the command save the predicted
probabilities alongside of the predicted segmentation masks requiring a lot of disk space.
Please select a separate OUTPUT_FOLDER
for each configuration!
Note that per default, inference will be done with all 5 folds from the cross-validation as an ensemble. We very
strongly recommend you use all 5 folds. Thus, all 5 folds must have been trained prior to running inference.
If you wish to make predictions with a single model, train the all
fold and specify it in nnUNetv2_predict
with -f all
If you wish to ensemble multiple predictions (typically form different configurations), you can do so with the following command:
nnUNetv2_ensemble -i FOLDER1 FOLDER2 ... -o OUTPUT_FOLDER -np NUM_PROCESSES
You can specify an arbitrary number of folders, but remember that each folder needs to contain npz files that were
generated by nnUNetv2_predict
. Again, nnUNetv2_ensemble -h
will tell you more about additional options.
Finally, apply the previously determined postprocessing to the (ensembled) predictions:
nnUNetv2_apply_postprocessing -i FOLDER_WITH_PREDICTIONS -o OUTPUT_FOLDER --pp_pkl_file POSTPROCESSING_FILE -plans_json PLANS_FILE -dataset_json DATASET_JSON_FILE
nnUNetv2_find_best_configuration
(or its generated inference_instructions.txt
file) will tell you where to find
the postprocessing file. If not you can just look for it in your results folder (it’s creatively named
postprocessing.pkl
). If your source folder is from an ensemble, you also need to specify a -plans_json
file and
a -dataset_json
file that should be used (for single configuration predictions these are automatically copied
from the respective training). You can pick these files from any of the ensemble members.
See here
Sometimes, the default 5-fold cross-validation split by nnU-Net does not fit a project. Maybe you want to run 3-fold
cross-validation instead? Or maybe your training cases cannot be split randomly and require careful stratification.
Fear not, for nnU-Net has got you covered (it really can do anything ❤️).
The splits nnU-Net uses are generated in the do_split
function of nnUNetTrainer. This function will first look for
existing splits, stored as a file, and if no split exists it will create one. So if you wish to influence the split,
manually creating a split file that will then be recognized and used is the way to go!
The split file is located in the nnUNet_preprocessed/DATASETXXX_NAME
folder. So it is best practice to first
populate this folder by running nnUNetv2_plan_and_preproccess
.
Splits are stored as a .json file. They are a simple python list. The length of that list is the number of splits it
contains (so it’s 5 in the default nnU-Net). Each list entry is a dictionary with keys ‘train’ and ‘val’. Values are
again simply lists with the train identifiers in each set. To illustrate this, I am just messing with the Dataset002
file as an example:
In [1]: from batchgenerators.utilities.file_and_folder_operations import load_json
In [2]: splits = load_json('splits_final.json')
In [3]: len(splits)
Out[3]: 5
In [4]: splits[0].keys()
Out[4]: dict_keys(['train', 'val'])
In [5]: len(splits[0]['train'])
Out[5]: 16
In [6]: len(splits[0]['val'])
Out[6]: 4
In [7]: print(splits[0])
{'train': ['la_003', 'la_004', 'la_005', 'la_009', 'la_010', 'la_011', 'la_014', 'la_017', 'la_018', 'la_019', 'la_020', 'la_022', 'la_023', 'la_026', 'la_029', 'la_030'],
'val': ['la_007', 'la_016', 'la_021', 'la_024']}
If you are still not sure what splits are supposed to look like, simply download some reference dataset from the
Medical Decathlon, start some training (to generate the splits) and manually inspect
the .json file with your text editor of choice!
In order to generate your custom splits, all you need to do is reproduce the data structure explained above and save it as
splits_final.json
in the nnUNet_preprocessed/DATASETXXX_NAME
folder. Then use nnUNetv2_train
etc. as usual.
So far nnU-Net only supports supervised pre-training, meaning that you train a regular nnU-Net on some source dataset
and then use the final network weights as initialization for your target dataset.
As a reminder, many training hyperparameters such as patch size and network topology differ between datasets as a
result of the automated dataset analysis and experiment planning nnU-Net is known for. So, out of the box, it is not
possible to simply take the network weights from some dataset and then reuse them for another.
Consequently, the plans need to be aligned between the two tasks. In this README we show how this can be achieved and
how the resulting weights can then be used for initialization.
Throughout this README we use the following terminology:
source dataset
is the dataset you intend to run the pretraining ontarget dataset
is the dataset you are interested in; the one you wish to fine tune onIn order to obtain matching network topologies we need to transfer the plans from one dataset to another. Since we are
only interested in the target dataset, we first need to run experiment planning (and preprocessing) for it:
nnUNetv2_plan_and_preprocess -d TARGET_DATASET
Then we need to extract the dataset fingerprint of the source dataset, if not yet available:
nnUNetv2_extract_fingerprint -d SOURCE_DATASET
Now we can take the plans from the target dataset and transfer it to the source:
nnUNetv2_move_plans_between_datasets -s SOURCE_DATSET -t TARGET_DATASET -sp SOURCE_PLANS_IDENTIFIER -tp TARGET_PLANS_IDENTIFIER
SOURCE_PLANS_IDENTIFIER
is hereby probably nnUNetPlans unless you changed the experiment planner in
nnUNetv2_plan_and_preprocess. For TARGET_PLANS_IDENTIFIER
we recommend you set something custom in order to not
overwrite default plans.
Note that EVERYTHING is transferred between the datasets. Not just the network topology, batch size and patch size but
also the normalization scheme! Therefore, a transfer between datasets that use different normalization schemes may not
work well (but it could, depending on the schemes!).
Note on CT normalization: Yes, also the clip values, mean and std are transferred!
Now you can run the preprocessing on the source task:
nnUNetv2_preprocess -d SOURCE_DATSET -plans_name TARGET_PLANS_IDENTIFIER
And run the training as usual:
nnUNetv2_train SOURCE_DATSET CONFIG all -p TARGET_PLANS_IDENTIFIER
Note how we use the ‘all’ fold to train on all available data. For pretraining it does not make sense to split the data.
Once pretraining is completed (or you obtain compatible weights by other means) you can use them to initialize your model:
nnUNetv2_train TARGET_DATASET CONFIG FOLD -pretrained_weights PATH_TO_CHECKPOINT
Specify the checkpoint in PATH_TO_CHECKPOINT.
When loading pretrained weights, all layers except the segmentation layers will be used!
So far there are no specific nnUNet trainers for fine tuning, so the current recommendation is to just use
nnUNetTrainer. You can however easily write your own trainers with learning rate ramp up, fine-tuning of segmentation
heads or shorter training time.