Does your system run like it should? Is your epoch time longer than expected? What epoch times should you expect?
Look no further for we have the solution here!
nnU-Net's benchmark trains models for 5 epochs. At the end, the fastest epoch will be noted down, along with the GPU name, torch version and cudnn version. You can find the benchmark output in the corresponding nnUNet_results subfolder (see example below). Don't worry, we also provide scripts to collect your results. Or you just start a benchmark and look at the console output. Everything is possible. Nothing is forbidden.
The benchmark implementation revolves around two trainers:
nnUNetTrainerBenchmark_5epochs
runs a regular training for 5 epochs. When completed, writes a .json file with the fastest epoch time as well as the GPU used and the torch and cudnn versions. Useful for speed testing the entire pipeline (data loading, augmentation, GPU training)nnUNetTrainerBenchmark_5epochs_noDataLoading
is the same, but it doesn't do any data loading or augmentation. It just presents dummy arrays to the GPU. Useful for checking pure GPU speed.
It's quite simple, actually. It looks just like a regular nnU-Net training.
We provide reference numbers for some of the Medical Segmentation Decathlon datasets because they are easily
accessible: download here. If it needs to be
quick and dirty, focus on Tasks 2 and 4. Download and extract the data and convert them to the nnU-Net format with
nnUNetv2_convert_MSD_dataset
.
Run nnUNetv2_plan_and_preprocess
for them.
Then, for each dataset, run the following commands (only one per GPU! Or one after the other):
nnUNetv2_train DATSET_ID 2d 0 -tr nnUNetTrainerBenchmark_5epochs
nnUNetv2_train DATSET_ID 3d_fullres 0 -tr nnUNetTrainerBenchmark_5epochs
nnUNetv2_train DATSET_ID 2d 0 -tr nnUNetTrainerBenchmark_5epochs_noDataLoading
nnUNetv2_train DATSET_ID 3d_fullres 0 -tr nnUNetTrainerBenchmark_5epochs_noDataLoading
If you want to inspect the outcome manually, check (for example!) your
nnUNet_results/DATASET_NAME/nnUNetTrainerBenchmark_5epochs__nnUNetPlans__3d_fullres/fold_0/
folder for the benchmark_result.json
file.
Note that there can be multiple entries in this file if the benchmark was run on different GPU types, torch versions or cudnn versions!
If you want to summarize your results like we did in our results, check the
summary script. Here you need to change the
torch version, cudnn version and dataset you want to summarize, then execute the script. You can find the exact
values you need to put there in one of your benchmark_result.json
files.
We have tested a variety of GPUs and summarized the results in a spreadsheet. Note that you can select the torch and cudnn versions at the bottom! There may be comments in this spreadsheet. Read them!
Results are shown as epoch time in seconds. Lower is better (duh). Epoch times can fluctuate between runs, so as long as you are within like 5-10% of the numbers we report, everything should be dandy.
If not, here is how you can try to find the culprit!
The first thing to do is to compare the performance between the nnUNetTrainerBenchmark_5epochs_noDataLoading
and
nnUNetTrainerBenchmark_5epochs
trainers. If the difference is about the same as we report in our spreadsheet, but
both your numbers are worse, the problem is with your GPU:
- Are you certain you compare the correct GPU? (duh)
- If yes, then you might want to install PyTorch in a different way. Never
pip install torch
! Go to the PyTorch installation page, select the most recent cuda version your system supports and only then copy and execute the correct command! Either pip or conda should work - If the problem is still not fixed, we recommend you try compiling pytorch from source. It's more difficult but that's how we roll here at the DKFZ (at least the cool kids here).
- Another thing to consider is to try exactly the same torch + cudnn version as we did in our spreadsheet. Sometimes newer versions can actually degrade performance and there might be bugs from time to time. Older versions are also often a lot slower!
- Finally, some very basic things that could impact your GPU performance:
- Is the GPU cooled adequately? Check the temperature with
nvidia-smi
. Hot GPUs throttle performance in order to not self-destruct - Is your OS using the GPU for displaying your desktop at the same time? If so then you can expect a performance penalty (I dunno like 10% !?). That's expected and OK.
- Are other users using the GPU as well?
- Is the GPU cooled adequately? Check the temperature with
If you see a large performance difference between nnUNetTrainerBenchmark_5epochs_noDataLoading
(fast) and
nnUNetTrainerBenchmark_5epochs
(slow) then the problem might be related to data loading and augmentation. As a
reminder, nnU-net does not use pre-augmented images (offline augmentation) but instead generates augmented training
samples on the fly during training (no, you cannot switch it to offline). This requires that your system can do partial
reads of the image files fast enough (SSD storage required!) and that your CPU is powerful enough to run the augmentations.
Check the following:
- [CPU bottleneck] How many CPU threads are running during the training? nnU-Net uses 12 processes for data augmentation by default.
If you see those 12 running constantly during training, consider increasing the number of processes used for data
augmentation (provided there is headroom on your CPU!). Increase the number until you see less active workers than
you configured (or just set the number to 32 and forget about it). You can do so by setting the
nnUNet_n_proc_DA
environment variable (Linux:export nnUNet_n_proc_DA=24
). Read here on how to do this. If your CPU does not support more processes (setting more processes than your CPU has threads makes no sense!) you are out of luck and in desperate need of a system upgrade! - [I/O bottleneck] If you don't see 12 (or nnUNet_n_proc_DA if you set it) processes running but your training times
are still slow then open up
top
(sorry, Windows users. I don't know how to do this on Windows) and look at the value left of 'wa' in the row that begins with '%Cpu (s)'. If this is >1.0 (arbitrarily set threshold here, essentially look for unusually high 'wa'. In a healthy training 'wa' will be almost 0) then your storage cannot keep up with data loading. Make sure to set nnUNet_preprocessed to a folder that is located on an SSD. nvme is preferred over SATA. PCIe3 is enough. 3000MB/s sequential read recommended. - [funky stuff] Sometimes there is funky stuff going on, especially when batch sizes are large, files are small and
patch sizes are small as well. As part of the data loading process, nnU-Net needs to open and close a file for each
training sample. Now imagine a dataset like Dataset004_Hippocampus where for the 2d config we have a batch size of
366 and we run 250 iterations in <10s on an A100. That's a lotta files per second (366 * 250 / 10 = 9150 files per second).
Oof. If the files are on some network drive (even if it's nvme) then (probably) good night. The good news: nnU-Net
has got you covered: add
export nnUNet_keep_files_open=True
to your .bashrc and the problem goes away. The neat part: it causes new problems if you are not allowed to have enough open files. You may have to increase the number of allowed open files.ulimit -n
gives your current limit (Linux only). It should not be something like 1024. Increasing that to 65535 works well for me. See here for how to change these limits: Link (works for Ubuntu 18, google for your OS!).