Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules

TL;DR

The paper tackles efficient discovery of data augmentation policies by learning augmentation schedules rather than fixed policies. It introduces Population Based Augmentation (PBA), which uses Population Based Training to jointly optimize augmentation parameters and training progress, producing a schedule that adapts across epochs. PBA achieves competitive results with AutoAugment on CIFAR-10/100 and SVHN while reducing compute by orders of magnitude, and ablations demonstrate that the schedule, not just the policy, drives gains. The approach is simple to implement within a PBT framework and is released as open source for broader adoption.

Abstract

A key challenge in leveraging data augmentation for neural network training is choosing an effective augmentation policy from a large search space of candidate operations. Properly chosen augmentation policies can lead to significant generalization improvements; however, state-of-the-art approaches such as AutoAugment are computationally infeasible to run for the ordinary user. In this paper, we introduce a new data augmentation algorithm, Population Based Augmentation (PBA), which generates nonstationary augmentation policy schedules instead of a fixed augmentation policy. We show that PBA can match the performance of AutoAugment on CIFAR-10, CIFAR-100, and SVHN, with three orders of magnitude less overall compute. On CIFAR-10 we achieve a mean test error of 1.46%, which is a slight improvement upon the current state-of-the-art. The code for PBA is open source and is available at https://github.com/arcelien/pba.

Figures (8)

• Figure 1: PBA matches AutoAugment's classification accuracy across a range of different network models on the CIFAR-10 dataset, while requiring 1,000x less GPU hours to run. For the full set of results, refer to Table \ref{['table-merged']}. Assuming an hourly GPU cost of $1.5, producing a new augmentation policy costs around $7.5 for PBA vs $7,500 with AutoAugment. The same scaling holds for the SVHN dataset as well.
• Figure 2: Comparison of AutoAugment and PBA augmentation strategies. In contrast to AutoAugment, PBA learns a schedule instead of a fixed policy. It does so in a short amount of time by using the PBT algorithm to jointly optimize augmentation policy parameters with the child model. PBA generates a single augmentation function $f(x, t)$ where $x$ is an input image and $t$ the current epoch, compared to AutoAugment's ensemble of augmentation policies $f^i(x)$, each of which has several further sub-policies.
• Figure 3: Augmentations applied to a CIFAR-10 "car" class image, at various points in our augmentation schedule learned on Reduced CIFAR-10 data. The maximum number of operations applied is sampled from 0 to 2. Each operation is formatted with name, probability, and magnitude value respectively.
• Figure 4: Plots showing the evolution of PBA operation parameters in the discovered schedule for CIFAR-10. Note that each operation actually appears in the parameter list twice; we take the mean parameter value for each operation in this visualization.
• Figure 5: Plot of the expected best child test accuracy after a given number of random trials on Wide-ResNet-40-2. Random policy schedules were generated by randomly selecting intervals of length between 1 and 40, and then selecting a random policy for the interval. All values were selected uniformly from the domain.