[fbsync] add prototype AugMix transform (#5492)

Summary:
* add prototype AugMix transform

* cleanup

* refactor auto augment subclasses to only trnasform a single image

* address review comments

Reviewed By: datumbox

Differential Revision: D34579510

fbshipit-source-id: b483642885913f844ffd72ef449de5a82467fc78

Author: prabhat00155

test/test_prototype_transforms.py

@@ -114,6 +114,7 @@ def test_mixup_cutmix(self, transform, input):
                 transforms.RandAugment(),
                 transforms.TrivialAugmentWide(),
                 transforms.AutoAugment(),
+                transforms.AugMix(),
             )
         ]
     )

torchvision/prototype/transforms/__init__.py

@@ -5,7 +5,7 @@
 from ._transform import Transform  # usort: skip
 
 from ._augment import RandomErasing, RandomMixup, RandomCutmix
-from ._auto_augment import RandAugment, TrivialAugmentWide, AutoAugment
+from ._auto_augment import RandAugment, TrivialAugmentWide, AutoAugment, AugMix
 from ._container import Compose, RandomApply, RandomChoice, RandomOrder
 from ._geometry import HorizontalFlip, Resize, CenterCrop, RandomResizedCrop
 from ._meta import ConvertBoundingBoxFormat, ConvertImageDtype, ConvertImageColorSpace