Performance of GPU-accelerated augmentations

[ziw-liu]

Most MONAI transforms that can be used for augmentations expect C(D)HW tensors without the batch dimension. This means that they have to be executed in python loops, and scales poorly on devices that benefit from parallel processing (GPUs or even just multi-threaded CPUs).

I did some benchmarking and found that MONAI's random affine transform can be more than 10x slower on the GPU than a batched implementation. Is there a way to accelerate MONAI's transforms? Or, in other words, am I doing something wrong in this benchmarking script?

Device	MONAI (loop)	Kornia (batched)	Relative
Zen 2 CPU (1 thread)	9160 ms	3800 ms	2.4x
Zen 2 CPU (16 threads)	7320 ms	556 ms	13.2x
A40 GPU	2620 ms	210 ms	12.5x

[KumoLiu]

Hi @ziw-liu, thank you for your feedback.
The transforms in MONAI follow the same design principles as PyTorch. They are designed for data processing in deep learning, specifically operating on individual data samples rather than batches. For batch processing, we recommend using DataLoader. I have also tested your benchmark using Dataloader, similar time usage is achieved. Kornia is primarily designed for general image processing tasks rather than deep learning workflows.
Hope it helps, thanks.

[ziw-liu]

Hi @KumoLiu and thanks for your suggestion!

We are currently using multi-processing-based workers in PyTorch DataLoader due to I/O and decoding that has to be done in Python. Is it safe/efficient to have each worker process execute transforms on the GPU independently?

[ziw-liu]

For context, PyTorch documentation states:

There are significant caveats to using CUDA models with multiprocessing; unless care is taken to meet the data handling requirements exactly, it is likely that your program will have incorrect or undefined behavior.