File size: 8,077 Bytes
ecf08bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
Changing the plans files grants you a lot of flexibility: You can depart from nnU-Net's default configuration and play
with different U-Net topologies, batch sizes and patch sizes. It is a powerful tool!
To better understand the components describing the network topology in our plans files, please read section 6.2
in the [supplementary information](https://static-content.springer.com/esm/art%3A10.1038%2Fs41592-020-01008-z/MediaObjects/41592_2020_1008_MOESM1_ESM.pdf)
(page 13) of our paper!
The goal of this tutorial is to demonstrate how to read and modify plans files and how to use them in your
experiments. The file used here works with Task120 and requires you to have downloaded the dataset, run
nnunet.dataset_conversion.Task120_Massachusetts_RoadSegm.py and then run nnUNet_plan_and_preprocess for it.
Note that this task is 2D only, but the same principles we use here can be easily extended to 3D and other tasks as well.
The output of `nnUNet_plan_and_preprocess` for this task looks like this:
[{'batch_size': 2,
'num_pool_per_axis': [8, 8],
'patch_size': array([1280, 1024]),
'median_patient_size_in_voxels': array([ 1, 1500, 1500]),
'current_spacing': array([999., 1., 1.]),
'original_spacing': array([999., 1., 1.]),
'pool_op_kernel_sizes': [[2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2]],
'conv_kernel_sizes': [[3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3]],
'do_dummy_2D_data_aug': False}]
This is also essentially what is saved in the plans file under the key 'plans_per_stage'
For this task, nnU-Net intends to use a patch size of 1280x1024 and a U-Net architecture with 8 pooling
operations per axis. Due to GPU memory constraints, the batch size is just 2.
Knowing the dataset we could hypothesize that a different approach might produce better results: The decision
of whether a pixel belongs to 'road' or not does not depend on the large contextual information that the large
patch size (and U-Net architecture) offer and could instead be made with more local information. Training with
a batch size of just 2 in a dataset with 800 training cases means that each batch contains only limited variability.
So one possible conclusion could be that smaller patches but larger batch sizes might result in a better
segmentation outcome. Let's investigate (using the same GPU memory constraint, determined manually with trial
and error!):
Variant 1: patch size 512x512, batch size 12
The following snippet makes the necessary adaptations to the plans file
```python
from batchgenerators.utilities.file_and_folder_operations import *
import numpy as np
from nnunet.paths import preprocessing_output_dir
task_name = 'Task120_MassRoadsSeg'
# if it breaks upon loading the plans file, make sure to run the Task120 dataset conversion and
# nnUNet_plan_and_preprocess first!
plans_fname = join(preprocessing_output_dir, task_name, 'nnUNetPlansv2.1_plans_2D.pkl')
plans = load_pickle(plans_fname)
plans['plans_per_stage'][0]['batch_size'] = 12
plans['plans_per_stage'][0]['patch_size'] = np.array((512, 512))
plans['plans_per_stage'][0]['num_pool_per_axis'] = [7, 7]
# because we changed the num_pool_per_axis, we need to change conv_kernel_sizes and pool_op_kernel_sizes as well!
plans['plans_per_stage'][0]['pool_op_kernel_sizes'] = [[2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2]]
plans['plans_per_stage'][0]['conv_kernel_sizes'] = [[3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3]]
# for a network with num_pool_per_axis [7,7] the correct length of pool kernel sizes is 7 and the length of conv
# kernel sizes is 8! Note that you can also change these numbers if you believe it makes sense. A pool kernel size
# of 1 will result in no pooling along that axis, a kernel size of 3 will reduce the size of the feature map
# representations by factor 3 instead of 2.
# save the plans under a new plans name. Note that the new plans file must end with _plans_2D.pkl!
save_pickle(plans, join(preprocessing_output_dir, task_name, 'nnUNetPlansv2.1_ps512_bs12_plans_2D.pkl'))
```
Variant 2: patch size 256x256, batch size 60
```python
from batchgenerators.utilities.file_and_folder_operations import *
import numpy as np
from nnunet.paths import preprocessing_output_dir
task_name = 'Task120_MassRoadsSeg'
plans_fname = join(preprocessing_output_dir, task_name, 'nnUNetPlansv2.1_plans_2D.pkl')
plans = load_pickle(plans_fname)
plans['plans_per_stage'][0]['batch_size'] = 60
plans['plans_per_stage'][0]['patch_size'] = np.array((256, 256))
plans['plans_per_stage'][0]['num_pool_per_axis'] = [6, 6]
plans['plans_per_stage'][0]['pool_op_kernel_sizes'] = [[2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2]]
plans['plans_per_stage'][0]['conv_kernel_sizes'] = [[3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3]]
save_pickle(plans, join(preprocessing_output_dir, task_name, 'nnUNetPlansv2.1_ps256_bs60_plans_2D.pkl'))
```
You can now use these custom plans files to train the networks and compare the results! Remeber that all nnUNet_*
commands have the -h argument to display their arguments. nnUNet_train supports custom plans via the -p argument.
Custom plans must be the prefix, so here this is everything except '_plans_2D.pkl':
Variant 1:
```bash
nnUNet_train 2d nnUNetTrainerV2 120 FOLD -p nnUNetPlansv2.1_ps512_bs12
```
Variant 2:
```bash
nnUNet_train 2d nnUNetTrainerV2 120 FOLD -p nnUNetPlansv2.1_ps256_bs60
```
Let all 5 folds run for each plans file (original and the two variants). To compare the results, you can make use of
nnUNet_determine_postprocessing to get the necessary metrics, for example:
```bash
nnUNet_determine_postprocessing -t 120 -tr nnUNetTrainerV2 -p nnUNetPlansv2.1_ps512_bs12 -m 2d
```
This will create a `cv_niftis_raw` and `cv_niftis_postprocessed` subfolder in the training output directory. In each
of these folders is a summary.json file that you can open with a regular text editor. In this file, there are metrics
for each training example in the dataset representing the outcome of the 5-fold cross-validation. At the very bottom
of the file, the metrics are aggregated through averaging (field "mean") and this is what you should be using to
compare the experiments. I recommend using the non-postprocessed summary.json (located in `cv_niftis_raw`) for this
because determining the postprocessing may actually overfit to the training dataset. Here are the results I obtained:
Vanilla nnU-Net: 0.7720\
Variant 1: 0.7724\
Variant 2: 0.7734
The results are remarkable similar and I would not necessarily conclude that such a small improvement in Dice is a
significant outcome. Nonetheless it was worth a try :-)
Despite the results shown here I would like to emphasize that modifying the plans file can be an extremely powerful
tool to improve the performance of nnU-Net on some datasets. You never know until you try it.
**ADDITIONAL INFORMATION (READ THIS!)**
- when working with 3d plans ('nnUNetPlansv2.1_plans_3D.pkl') the 3d_lowres and 3d_fullres stage will be encoded
in the same plans file. If len(plans['plans_per_stage']) == 2, then [0] is the 3d_lowres and [1] is the
3d_fullres variant. If len(plans['plans_per_stage']) == 1 then [0] will be 3d_fullres and 3d_cascade_fullres
(they use the same plans).
- 'pool_op_kernel_sizes' together with determines 'patch_size' determines the size of the feature map
representations at the bottleneck. For Variant 1 & 2 presented here, the size of the feature map representation is
`print(plans['plans_per_stage'][0]['patch_size'] / np.prod(plans['plans_per_stage'][0]['pool_op_kernel_sizes'], 0))`
> [4., 4.]
If you see a non-integer number here, your model will crash! Make sure these are always integers!
nnU-Net will never create smaller bottlenecks than 4!
- do not change the 'current_spacing' in the plans file! This will not work properly. To change the target spacing,
have a look at the [custom spacing](custom_spacing.md) tutorial. |