Batched (random) augmentations

🚀 The feature

Randomly sample augmentation parameters for each image in the batch

Motivation, pitch

Torch augmentations have two big advantages over alternatives like Albumentations: ability to apply them to a batch of images and ability to run them on GPU. This can make a big difference when batch size is very large (e.g., 1000). However, right now augmentation pipeline would apply exactly the same transformation to every image in the batch, which makes augmentations less valuable and can destabilize the training because of significantly chaged batch statistics.

Alternatives

Any recommendations to existing libraries that allow batched + GPU augmentations?

Additional context

In some of my training pipelines I observed that Albumentations are not fast enough and that we can benefit from going back to torchvision, but then I noticed that all images are augmented in the same way.

The issue does not name files, tests, or an implementation entry point. Start by locating torchvision's batched augmentation pipeline and existing parameter-sampling behavior; done would mean defining and validating per-image random parameters while preserving batched and GPU execution.