Faster AutoAugment: Learning Augmentation Strategies using Backpropagation

If you work with image classifiers, chances are that you have augmented your data. Hand-crafted pipelines were the norm for a while until automated search over augmentation choices took the field by storm around 2018 (if you had the computational resources for it). [Cub19A], which is perhaps amongst the most popular ones, discretises the whole search space of image transformations, including their parameters, and uses policy gradient to select “augmentation policies” based on the validation score of the target neural network over an image batch.

Alas, policy gradient (and other black-box approaches like Bayesian optimization) can be extremely costly. Inspired by ideas from differentiable neural architecture search, [Hat20F] offers an alternative which can be significantly faster while achieving similar performance

Replacing search and sampling of categorical variables by differentiable operations during training, falling back to discrete operations at inference time. They achieve this by modelling the search as a differentiable process. First, sampling of a Bernoulli variable for deciding whether to apply a transformation is replaced by sampling from a relaxed continuous distribution, building on [Jan17C]. Second, they provide approximations to the gradients of several non-differentiable image transformations and re-implement differentiable transformations in pytorch (as opposed to using PIL like AutoAugment does). While during training the sampling of discrete quantities is replaced by these differentiable operations, at inference time the method falls back to discrete operations.

Finally, they do not directly optimize final model performance and instead minimise the distance of the augmented images to the original distribution to encourage “filling the holes in the dataset”. A classification loss is added as well to avoid flipping labels.

There are of course lots of details under the hood, like Wasserstein GANs, network architecture choices and DARTS. Follow the links in the bibliography to read the full paper or experiment with the code.

Data augmentation is an effective technique for improving the accuracy of modern image classifiers. However, current data augmentation implementations are manually designed. In this paper, we describe a simple procedure called AutoAugment to automatically search for improved data augmentation policies. In our implementation, we have designed a search space where a policy consists of many …

Data augmentation methods are indispensable heuristics to boost the performance of deep neural networks, especially in image recognition tasks. Recently, several studies have shown that augmentation strategies found by search algorithms outperform hand-made strategies. Such methods employ black-box search algorithms over image transformations with continuous or discrete parameters and require a …

Categorical variables are a natural choice for representing discrete structure in the world. However, stochastic neural networks rarely use categorical latent variables due to the inability to backpropagate through samples. In this work, we present an efficient gradient estimator that replaces the non-differentiable sample from a categorical distribution with a differentiable sample from a novel …

References