Image Augmentations on GPU Tests

Hello Pytorch vision people !

I am currently working on a project that requires lots of image augmentations
to perform better. And I believe this is not only my case. When reading
about topics such a domain randomization, we see that big variations on images
leads to much more generalization.

I saw that pytorch does not seem to provide a way to perform
any image augmentation on GPU as comment in #45 . In some posts I saw people
not encouraging to do it (https://discuss.pytorch.org/t/preprocess-images-on-gpu/5096) but i really disagree, specially for the cases where several augmentations are applied.
To show this point I provide a gist code showing an example illustrating the
possible speed up gains on a multiplication operation ( brightness augmentation ? )

https://gist.github.com/felipecode/f3531e2d04e846da99053aff16b06028

On the gist, i show a GPU augmentation interface is working as following:

no_aug_trans = transforms.Compose( [transforms.RandomResizedCrop(224), transforms.ToTensor()])  
dataset = datasets.ImageFolder(data_path, transform=no_aug_trans)
multply_gpu = transforms.Compose([ToGPU()] + [Multiply(1.01)] * number_of_multiplications)
for data in data_loader:  
    image, labels = data  
    result = multiply_gpu(image)

Unfortunately the GPU augmentation could not be smoothly interfaced with the dataloader without
sacrificing the multi threading for data reading. However, the speed ups obtained seems promising
The following plot shows up when running the gist code with a TITAN Xp and Intel(R) Xeon(R) CPU E5-1620 v3 @ 3.50GHz CPU, note that I remove the loading time when plotting.

The plot shows the time to compute in function of number of multiplications. For this test, on each data point about 500 RGB images of 224x224 are multiplied by a constant.

Of course, there is no clear reason on why should one do 60 multiplications.
However, I implemented an small library where I used imgaug library as reference
and implemented more functions in GPU. For the following augmentation set
used in my project I obtained about 3-4 times speed up.

transforms.Compose( [ToGPU(), Add((-5, 5)), Multiply((0.9, 1.1)), Dropout(0.2),AdditiveGaussianNoise(0.10*255),GaussianBlur(sigma=(0.0, 3.0)),ContrastNormalization((0.5, 1.5))] )

This speed up is even higher if more augmentations are added.

So, how can I improve this API ? How could something like this fit in a pull request ?
How can this be more smoothly merged inside the dataloader , but keeping the
multithreading for data reading ?
I still have to test the training time for the full system, but I don't believe there will
be any overhead since images have to be copyed to GPU anyway.

Start with the linked gist and the transforms.Compose, datasets.ImageFolder, and data_loader entry points described in the issue. Review how existing transforms and data loading handle device placement and multiprocessing. The issue needs a settled API and integration scope before completion can be defined.