This PR adds the experimental nnUNetTrainerV2Weighted trainer class. It also adds the following to the docs:
WARNING this trainer is still very experimental, also the training times can increase significantly using this trainer, w.r.t. other trainers!
This trainer class implements a simple training scheme with support for training with weightmaps (annotated weight values at the voxel/pixel level).
Note that for this trainer to work weightmaps have to be supplied next to the preprocessed training data.
First you would prepare your data as normal and run all the planning and preprocessing steps.
Once you have all your preprocessed files inside the nnUNet_preprocessed_data, you put the weightmaps next to the preprocessed files as .npy files here: nnUNet_preprocessed_data/TaskXXX_DIRECTORY/nnUNetData_plans_v2.1_stage0.
So that you have the files similar to the following formats there:
To simplify the weightmap generation processes we have added a utility CLI that can be used to generate weightmaps in a similar fashion to the original U-Net paper:
nnUNet_diag_generate_unet_weightmaps --help
Some caveats with the weightmap training, which are good to know before usage:
You can experience very long training times. The data pipeline is currently not optimized to handle the weightmaps and takes a lot of additional overhead.
2D training is currently not supported. If you need it let me know, it should be easy to add.
Dice loss has been disabled by default during training, since it does not support training with weightmaps.
The reported Dice values during online evaluation do not take weightmaps into account and should be ignored.
Closes #7
This PR adds the experimental
nnUNetTrainerV2Weighted
trainer class. It also adds the following to the docs:WARNING this trainer is still very experimental, also the training times can increase significantly using this trainer, w.r.t. other trainers!
This trainer class implements a simple training scheme with support for training with weightmaps (annotated weight values at the voxel/pixel level).
Note that for this trainer to work weightmaps have to be supplied next to the preprocessed training data. First you would prepare your data as normal and run all the planning and preprocessing steps. Once you have all your preprocessed files inside the
nnUNet_preprocessed_data
, you put the weightmaps next to the preprocessed files as.npy
files here:nnUNet_preprocessed_data/TaskXXX_DIRECTORY/nnUNetData_plans_v2.1_stage0
. So that you have the files similar to the following formats there:To simplify the weightmap generation processes we have added a utility CLI that can be used to generate weightmaps in a similar fashion to the original U-Net paper:
nnUNet_diag_generate_unet_weightmaps --help
Some caveats with the weightmap training, which are good to know before usage: