讨论: 批量 GPU 数据增强

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 spent on 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. 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