Hi! For a PhD course I profiled the nnU-Net training loop with NVTX + Nsight Systems (1x A100, `3d_fullres`, batch size 2, preprocessed data on an NVMe partition). With 8 CPUs, the GPU was starving and most of each iteration was spent waiting on the CPU workers, and inside the workers most of the time was the augmentations.
I tried moving the augmentations to the GPU and running them on the whole batch at once, as an opt-in trainer variant, motivated by [batchaug](https://github.com/halleewong/batchaug/tree/main). Workers only load, crop/pad and stack the raw data (`transforms=None`), and the transform pipeline runs in `train_step` right after the batch is moved to the device. It roughly halves the iteration time, and the speedup persisted even with 4 CPUs. **Before opening a PR, I wanted to ask whether this is a direction you'd want in nnU-Net at all.**
## Numbers
Measurement | CPU transforms (8 workers) | GPU batched transforms (8 workers)
-- | -- | --
GPU utilization | 39.9% | 86.0%
CPU utilization | 94.5% | 39.9%
iteration | 302.4 ms | 145.5 ms
`next()` on dataloader | 173.0 ms | 0.42 ms
train step | 130.7 ms | 145.0 ms
transforms | 1022.9 ms/sample (CPU) | 13.5 ms/batch (GPU)
epoch time | ~70 s | ~39 s
The augmentations go from ~1 s per sample to ~13 ms per batch, and the wait on the dataloader basically disappears.
train step grows by ~14 ms, due to the GPU transforms now happening inside it.
Two other things I found interesting:
- I tested it with 4 CPUs and workers. With CPU augmentations, going 8 -> 4 workers nearly doubles the iteration time (302 -> 577 ms) and GPU utilization drops to 20%. With GPU augmentations, 8 -> 4 workers changes almost nothing (145 -> 151 ms) and GPU utilization stays ~84%. Useful on nodes with few cores per GPU.
- It could help when training with large batch size. #2361 reports slow-downs when increasing batch size, which is understandable if transforms are done per-sample on the CPU (you would need more workers to keep the GPU fed at the same rate). Moving the transforms to the GPU should leave to the CPUs only the loading and decompression (Blosc), vastly reducing this phenomenon.
Screenshots of the two timelines (yellow block is `next()` on dataloader, blue block is forward + backward, red block is yellow + blue, orange block is data loading from disk, cyan block is the transforms):
CPU transforms, 8 workers
 |
 |
| GPU view |
CPU view |
GPU transforms, 8 workers
 |
 |
| GPU view |
CPU view |
## Does it still train the same?
I trained the same config twice on the PanTS dataset (abdominal CT), once with the stock `nnUNetTrainer` and once with the GPU version. Attaching both `summary.json`, the numbers are pretty similar. I don't have time however to evaluate on more datasets unfortunately.
[summary_cpu.json](https://github.com/user-attachments/files/31906846/summary_cpu.json)
[summary_batchaug.json](https://github.com/user-attachments/files/31906845/summary_batchaug.json)
## What the implementation looks like
- [batchaug](https://github.com/halleewong/batchaug) for the transforms that already have a batched GPU equivalent (noise, blur, low resolution, mirroring);
- Claude implemented custom GPU versions for the remaining transforms: gamma (`invert_image` / `retain_stats`), contrast and multiplicative brightness (the split-range sampling + `preserve_range`), rotation+scaling, and batched `RemoveLabel` / `ConvertSegmentationToRegions` / `DownsampleSegForDS`. I run tests to check that the two versions matched.
I didn't have time to derive them myself just for a small course project. The profiling though shows clear benefits and I think the direction is worth it, so I'd understand if liked the idea but decided to reimplement them yourselves.
## One more reason I think this will matter
C-Blosc2 4.0 plans to move decompression to the GPU (From their [4.0 roadmap](https://blosc.org/c-blosc2/development/roadmap.html): "The idea is to offload the compression, but most importantly, decompression tasks to the GPU, so that the CPU is free to do other tasks."). Using GPU decompression in nnUNet would be impossible if augmentations stay on the CPU.
Happy to open a PR with the trainer variant + tests or to just leave this here as a profiling report if you're not interested. Let me know! :)