I am currently using the RandAugment class from tf-models (from official.vision.beta.ops import augment). The RandAugment().distort(), however, does not allow batched inputs, and computation-wise it’s expensive as well (especially when you have more than two augmentation operations).
So, following suggestions from this guide 4, I wanted to be able to map RandAugment().distort() after my dataset is batched. Any workaround for that?
Here’s how I am building my input pipeline for now:
# Recommended is m=2, n=9
augmenter = augment.RandAugment(num_layers=3, magnitude=10)

dataset = load_dataset(filenames)
dataset = dataset.shuffle(batch_size*10)
dataset = dataset.map(augmenter.distort, num_parallel_calls=AUTO)