Modifying the nnU-Net Configurations
nnU-Net provides unprecedented out-of-the-box segmentation performance for essentially any dataset we have evaluated
it on. That said, there is always room for improvements. A fool-proof strategy for squeezing out the last bit of
performance is to start with the default nnU-Net, and then further tune it manually to a concrete dataset at hand.
This guide is about changes to the nnU-Net configuration you can make via the plans files. It does not cover code
extensions of nnU-Net. For that, take a look here
In nnU-Net V2, plans files are SO MUCH MORE powerful than they were in v1. There are a lot more knobs that you can
turn without resorting to hacky solutions or even having to touch the nnU-Net code at all! And as an added bonus:
plans files are now also .json files and no longer require users to fiddle with pickle. Just open them in your text
editor of choice!
If overwhelmed, look at our Examples!
plans.json structure
Plans have global and local settings. Global settings are applied to all configurations in that plans file while
local settings are attached to a specific configuration.
Global settings
foreground_intensity_properties_by_modality
: Intensity statistics of the foreground regions (all labels except
background and ignore label), computed over all training cases. Used by CT normalization scheme.
image_reader_writer
: Name of the image reader/writer class that should be used with this dataset. You might want
to change this if, for example, you would like to run inference with files that have a different file format. The
class that is named here must be located in nnunetv2.imageio!
label_manager
: The name of the class that does label handling. Take a look at
nnunetv2.utilities.label_handling.LabelManager to see what it does. If you decide to change it, place your version
in nnunetv2.utilities.label_handling!
transpose_forward
: nnU-Net transposes the input data so that the axes with the highest resolution (lowest spacing)
come last. This is because the 2D U-Net operates on the trailing dimensions (more efficient slicing due to internal
memory layout of arrays). Future work might move this setting to affect only individual configurations.
- transpose_backward is what numpy.transpose gets as new axis ordering.
transpose_backward
: the axis ordering that inverts "transpose_forward"
- [
original_median_shape_after_transp
]: just here for your information
- [
original_median_spacing_after_transp
]: just here for your information
- [
plans_name
]: do not change. Used internally
- [
experiment_planner_used
]: just here as metadata so that we know what planner originally generated this file
- [
dataset_name
]: do not change. This is the dataset these plans are intended for
Local settings
Plans also have a configurations
key in which the actual configurations are stored. configurations
are again a
dictionary, where the keys are the configuration names and the values are the local settings for each configuration.
To better understand the components describing the network topology in our plans files, please read section 6.2
in the supplementary information
(page 13) of our paper!
Local settings:
spacing
: the target spacing used in this configuration
patch_size
: the patch size used for training this configuration
data_identifier
: the preprocessed data for this configuration will be saved in
nnUNet_preprocessed/DATASET_NAME/data_identifier. If you add a new configuration, remember to set a unique
data_identifier in order to not create conflicts with other configurations (unless you plan to reuse the data from
another configuration, for example as is done in the cascade)
batch_size
: batch size used for training
batch_dice
: whether to use batch dice (pretend all samples in the batch are one image, compute dice loss over that)
or not (each sample in the batch is a separate image, compute dice loss for each sample and average over samples)
preprocessor_name
: Name of the preprocessor class used for running preprocessing. Class must be located in
nnunetv2.preprocessing.preprocessors
use_mask_for_norm
: whether to use the nonzero mask for normalization or not (relevant for BraTS and the like,
probably False for all other datasets). Interacts with ImageNormalization class
normalization_schemes
: mapping of channel identifier to ImageNormalization class name. ImageNormalization
classes must be located in nnunetv2.preprocessing.normalization. Also see here
resampling_fn_data
: name of resampling function to be used for resizing image data. resampling function must be
callable(data, current_spacing, new_spacing, **kwargs). It must be located in nnunetv2.preprocessing.resampling
resampling_fn_data_kwargs
: kwargs for resampling_fn_data
resampling_fn_probabilities
: name of resampling function to be used for resizing predicted class probabilities/logits.
resampling function must be callable(data, current_spacing, new_spacing, **kwargs). It must be located in
nnunetv2.preprocessing.resampling
resampling_fn_probabilities_kwargs
: kwargs for resampling_fn_probabilities
resampling_fn_seg
: name of resampling function to be used for resizing segmentation maps (integer: 0, 1, 2, 3, etc).
resampling function must be callable(data, current_spacing, new_spacing, **kwargs). It must be located in
nnunetv2.preprocessing.resampling
resampling_fn_seg_kwargs
: kwargs for resampling_fn_seg
UNet_class_name
: UNet class name, can be used to integrate custom dynamic architectures
UNet_base_num_features
: The number of starting features for the UNet architecture. Default is 32. Default: Features
are doubled with each downsampling
unet_max_num_features
: Maximum number of features (default: capped at 320 for 3D and 512 for 2d). The purpose is to
prevent parameters from exploding too much.
conv_kernel_sizes
: the convolutional kernel sizes used by nnU-Net in each stage of the encoder. The decoder
mirrors the encoder and is therefore not explicitly listed here! The list is as long as n_conv_per_stage_encoder
has
entries
n_conv_per_stage_encoder
: number of convolutions used per stage (=at a feature map resolution in the encoder) in the encoder.
Default is 2. The list has as many entries as the encoder has stages
n_conv_per_stage_decoder
: number of convolutions used per stage in the decoder. Also see n_conv_per_stage_encoder
num_pool_per_axis
: number of times each of the spatial axes is pooled in the network. Needed to know how to pad
image sizes during inference (num_pool = 5 means input must be divisible by 2**5=32)
pool_op_kernel_sizes
: the pooling kernel sizes (and at the same time strides) for each stage of the encoder
- [
median_image_size_in_voxels
]: the median size of the images of the training set at the current target spacing.
Do not modify this as this is not used. It is just here for your information.
Special local settings:
inherits_from
: configurations can inherit from each other. This makes it easy to add new configurations that only
differ in a few local settings from another. If using this, remember to set a new data_identifier
(if needed)!
previous_stage
: if this configuration is part of a cascade, we need to know what the previous stage (for example
the low resolution configuration) was. This needs to be specified here.
next_stage
: if this configuration is part of a cascade, we need to know what possible subsequent stages are! This
is because we need to export predictions in the correct spacing when running the validation. next_stage
can either
be a string or a list of strings
Examples
Increasing the batch size for large datasets
If your dataset is large the training can benefit from larger batch_sizes. To do this, simply create a new
configuration in the configurations
dict
"configurations": {
"3d_fullres_bs40": {
"inherits_from": "3d_fullres",
"batch_size": 40
}
}
No need to change the data_identifier. 3d_fullres_bs40
will just use the preprocessed data from 3d_fullres
.
No need to rerun nnUNetv2_preprocess
because we can use already existing data (if available) from 3d_fullres
.
Using custom preprocessors
If you would like to use a different preprocessor class then this can be specified as follows:
"configurations": {
"3d_fullres_my_preprocesor": {
"inherits_from": "3d_fullres",
"preprocessor_name": MY_PREPROCESSOR,
"data_identifier": "3d_fullres_my_preprocesor"
}
}
You need to run preprocessing for this new configuration:
nnUNetv2_preprocess -d DATASET_ID -c 3d_fullres_my_preprocesor
because it changes the preprocessing. Remember to
set a unique data_identifier
whenever you make modifications to the preprocessed data!
Change target spacing
"configurations": {
"3d_fullres_my_spacing": {
"inherits_from": "3d_fullres",
"spacing": [X, Y, Z],
"data_identifier": "3d_fullres_my_spacing"
}
}
You need to run preprocessing for this new configuration:
nnUNetv2_preprocess -d DATASET_ID -c 3d_fullres_my_spacing
because it changes the preprocessing. Remember to
set a unique data_identifier
whenever you make modifications to the preprocessed data!
Adding a cascade to a dataset where it does not exist
Hippocampus is small. It doesn't have a cascade. It also doesn't really make sense to add a cascade here but hey for
the sake of demonstration we can do that.
We change the following things here:
spacing
: The lowres stage should operate at a lower resolution
- we modify the
median_image_size_in_voxels
entry as a guide for what original image sizes we deal with
- we set some patch size that is inspired by
median_image_size_in_voxels
- we need to remember that the patch size must be divisible by 2**num_pool in each axis!
- network parameters such as kernel sizes, pooling operations are changed accordingly
- we need to specify the name of the next stage
- we need to add the highres stage
This is how this would look like (comparisons with 3d_fullres given as reference):
"configurations": {
"3d_lowres": {
"inherits_from": "3d_fullres",
"data_identifier": "3d_lowres"
"spacing": [2.0, 2.0, 2.0], # from [1.0, 1.0, 1.0] in 3d_fullres
"median_image_size_in_voxels": [18, 25, 18], # from [36, 50, 35]
"patch_size": [20, 28, 20], # from [40, 56, 40]
"n_conv_per_stage_encoder": [2, 2, 2], # one less entry than 3d_fullres ([2, 2, 2, 2])
"n_conv_per_stage_decoder": [2, 2], # one less entry than 3d_fullres
"num_pool_per_axis": [2, 2, 2], # one less pooling than 3d_fullres in each dimension (3d_fullres: [3, 3, 3])
"pool_op_kernel_sizes": [[1, 1, 1], [2, 2, 2], [2, 2, 2]], # one less [2, 2, 2]
"conv_kernel_sizes": [[3, 3, 3], [3, 3, 3], [3, 3, 3]], # one less [3, 3, 3]
"next_stage": "3d_cascade_fullres" # name of the next stage in the cascade
},
"3d_cascade_fullres": { # does not need a data_identifier because we can use the data of 3d_fullres
"inherits_from": "3d_fullres",
"previous_stage": "3d_lowres" # name of the previous stage
}
}
To better understand the components describing the network topology in our plans files, please read section 6.2
in the supplementary information
(page 13) of our paper!