DIAGNijmegen / nnUNet

Apache License 2.0
3 stars 14 forks source link

Add custom trainer for training with weightmaps #19

Closed silvandeleemput closed 2 years ago

silvandeleemput commented 2 years ago

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:

├── hippocampus_001.npy
├── hippocampus_001.npz
├── hippocampus_001.pkl
├── hippocampus_001_weightmap.npy
├── hippocampus_002.npy
├── hippocampus_002.npz
├── hippocampus_002.pkl
└── hippocampus_002_weightmap.npy

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:

silvandeleemput commented 2 years ago

Nice, looks good! Did you try a full training run?

I tried a full training run and it works. However, it is very slow w.r.t. training without weightmaps.