Safe Augmentation: Learning Task-Specific Transformations from Data

TL;DR

Safe Augmentation is a simple, explainable method that learns task-specific data transformations to improve deep learning model generalization without altering data distribution, outperforming manual augmentation techniques.

Contribution

The paper introduces Safe Augmentation, a novel approach to automatically learn task-specific, distribution-preserving data transformations for better model generalization.

Findings

Improves accuracy on multiple datasets including CIFAR-10, CIFAR-100, SVHN, Tiny ImageNet, and Cityscapes.

Model-agnostic and easy to implement.

Outperforms baseline augmentation methods.

Abstract

Data augmentation is widely used as a part of the training process applied to deep learning models, especially in the computer vision domain. Currently, common data augmentation techniques are designed manually. Therefore they require expert knowledge and time. Moreover, augmentations are dataset-specific, and the optimal augmentations set on a specific dataset has limited transferability to others. We present a simple and explainable method called $\textbf{Safe Augmentation}$ that can learn task-specific data augmentation techniques that do not change the data distribution and improve the generalization of the model. We propose to use safe augmentation in two ways: for model fine-tuning and along with other augmentation techniques. Our method is model-agnostic, easy to implement, and achieves better accuracy on CIFAR-10, CIFAR-100, SVHN, Tiny ImageNet, and Cityscapes datasets comparing to baseline augmentation techniques. The code is available at $\href{https://github.com/Irynei/SafeAugmentation}{https://github.com/Irynei/SafeAugmentation}$.

Tables

Table 1. Table 1: Safe augmentations found using joint learning setup for CIFAR-10, CIFAR-100, Tiny ImageNet and SVHN datasets.

	C-10	C-100	Tiny INet	SVHN
HorizontalFlip	✓	✓	✓
VerticalFlip
RandomRotate90	✓	✓	✓
Transpose
ToGray
ShiftScaleRotate
RandomCrop	✓	✓	✓
CenterCrop	✓	✓	✓
RandomSizedCrop	✓	✓	✓
RandomContrast		✓		✓
RandomBrightness	✓	✓	✓	✓
RandomGamma	✓	✓	✓	✓
CLAHE
Blur
GaussNoise				✓

Table 2. Table 2: Test top-1 accuracy (%). All results are averaged over 3 runs. All fine-tuned experiments were performed using models pre-trained on all augmentations. * Safe v2 is defined manually by removing RandomCrop and CenterCrop.

	CIFAR-10		SVHN		CIFAR-100	Tiny ImageNet
	DenseNet 121	DenseNet 169	DenseNet 121	DenseNet 169	DenseNet 121	DenseNet 121
Without	79.39	79.52	95.87	96.19	53.55	49.60
Baseline	87.31	87.15	95.99	95.63	61.67	49.73
Safe	87.85	86.84	96.14	96.67	64.17	57.61
All	87.79	88.15	96.19	96.43	65.82	58.66
Fine-tuned on All	88.68	88.15	96.19	96.58	65.83	58.65
Fine-tuned on Safe	88.38	88.21	96.36	97.01	65.93	58.85
Fine-tuned on Safe v2*	88.59	88.46	-	-	65.99	59.00

Table 3. Table 3: Safe augmentations found using joint learning setup for the Cityscapes [ 5 ] dataset.

Augmentations

HorizontalFlip

RandomBrightness

RandomGamma

Transpose

Table 4. Table 4: Validation IoU(%) on the Cityscapes [ 5 ] dataset on single FPN model with DenseNet-121 backbone

	FPN(DenseNet-121)
Without	45.34
Baseline	51.11
Safe	51.58
All	59.41
Fine-tuned on All	60.37
Fine-tuned on Safe	62.09

Table 5. Table 5: Test top-1 accuracy (%). Comparison with Cutout [ 9 ] and AutoAugment [ 6 ] . Safe* means Safe for SVHN and Safe v2 for CIFAR-10 and CIFAR-100

	CIFAR-10		SVHN		CIFAR-100
	DenseNet 121	DenseNet 169	DenseNet 121	DenseNet 169	DenseNet 121
Without	79.39	79.52	95.87	96.19	53.55
Best Policy AA	84.60	85.51	96.67	96.65	59.34
Baseline	87.31	87.15	95.99	95.63	61.67
Baseline + Cutout	88.10	88.84	96.31	96.39	63.26
\rowcolor[HTML]E5F7E5 Finetuned Safe*	88.59	88.46	96.36	97.01	65.99
\rowcolor[HTML]E5F7E5 Safe* + Baseline + Cutout	88.16	88.32	96.73	96.39	65.39
AutoAugment	90.77	90.58	96.76	96.66	68.70

Equations

$$L_{total}=L_{augm}+L_{class}$$

$$L_{total}=L_{augm}+L_{segm}$$

Full text

1 Introduction

Deep neural networks achieve human-level or even higher performance in many computer vision tasks, such as image classification, image restoration, image or video segmentation, etc. [34]. For example, the human top-5 image classification error on the ImageNet dataset is $5\%$, whereas the current state-of-the-art deep neural networks achieve nearly $3\%$ [28].

However, deep learning models require a massive amount of training data to be robust. Data augmentation is one of the approaches that can help to handle this issue by expanding training data using transformations that preserve semantic information and class labels. The choice of the data augmentation techniques to use for the specific dataset and task is not a trivial one. While some augmentations increase the performance and generalization of the model, others can have a negative impact. For instance, a horizontal flip is proven to be useful augmentation for ImageNet-like datasets, but not for the MNIST dataset [20], as it changes the distribution of the data because horizontally flipped digits are often no longer valid digits.

Automated machine learning, in particular, automatic data augmentation is currently am important research topic. AutoAugment uses reinforcement learning to search for optimal augmentation policies along with the magnitudes and probabilities [6]. While these methods provide superior image classification results, they often have limited explainability and require lots of computational resources. Yet the majority of common data augmentation techniques are designed either empirically or by leveraging expert knowledge. Hence, it decreases the transferability of such methods between different tasks.

This paper aims to make another push on making the process of choosing data augmentations automatic and explainable, namely to learn from data which augmentation techniques lead to model generalization improvement. Our contributions are summarized below.

• •

We introduce a simple, intuitive model-agnostic method for choosing augmentations that can be safely used during model training. In our implementation, we take a fixed set of the common image processing functions with default magnitudes and select a subset of augmentations which produce images that comes from the same distribution that existing ones while improving the accuracy of the main task(e.g., image classification, image segmentation).

• •

We propose two ways of using learned augmentations: for model fine-tuning and along with other augmentation techniques. Our experiments on different datasets and tasks show that Safe Augmentation works better than baseline augmentations and comparable with more advanced augmentation techniques while being intuitive, explainable, and straightforward.

2 Related Work

2.1 Data Augmentation

The recent paper by Hernández-García and König has shown that data augmentation alone can achieve the same or even higher performance than explicit regularization techniques (weight decay [19], dropout [29], etc.), without wasting model capacity [12].

Traditional augmentation. The most common type of data augmentations are geometric transformations, such as flipping, cropping, rotating, scaling, etc., and color transformations, such as adjusting color, brightness, resolution, etc. They are often called generic or traditional augmentations. They all fall under the category of data warping and are usually performed in the data space, e.g., Wong et al. have shown that it is more efficient to perform data augmentation in data space than in feature space as long as label preserving transforms are known [32]. This type of transformations is easy to use and efficient to implement. One main disadvantage is that you need to have expert knowledge in the image domain to choose transformations that will not affect the correctness of the image labels. Traditional augmentations are broadly used and have shown excellent results in reducing overfitting and improving model performance [26] [30].

Generative Adversarial Networks (GANs). In 2014 Goodfellow et al. proposed a new class of neural networks that can generate realistic data from scratch using generator and discriminator networks that are trained in the minimax two-player game framework [10]. GANs can be used as a form of unsupervised data augmentation by generating new data from the source distribution. They have also been used for style transfer, e.g., transferring images from one weather condition to another. These generated images can be used to help the model to work in different conditions, for instance, to train autonomous cars to drive in night or snow, having collected data from sunny weather only. GANs are also shown to be successfully used for data augmentation in the medical imaging domain by synthetically augment mammogram and MRI images [3, 33].

2.2 Automated Data Augmentation

Despite described advantages, common data augmentation methods are usually dataset-specific and designed manually, which require prior expert knowledge and time. Recently, many researches were focused on the automation of the process of data augmentation. We divide them into the following two groups.

Generate augmented data directly. Smart Augmentation proposed by Lemley, Bazrafkan, and Corcoran can automatically generate augmented data by merging two or more samples from the same class, in a way that reduces the loss of the original model [21]. DeVries and Taylor proposed a domain-independent data augmentation technique by using simple transformations in the learned feature space. They train a sequence autoencoder to construct a learned feature space in which they extrapolate between samples [8]. Tran et al. introduced a novel Bayesian method for generating additional data based on the distribution learned from the training set [31]. Generative adversarial networks have been extensively used for producing augmented data. For example, Antoniou, Storkey, and Edwards presented DAGAN - An image conditional GAN-based model that learns from one data item how to generate other realistic within-class data items. DAGAN can be applied to unseen classes of data and can also enhance few-shot learning systems. [2]. Another approach called DADA: Deep Adversarial Data Augmentation was proposed by Zhang et al. to train deep learning models in extremely low data regimes. They show that that DADA outperforms both traditional data augmentation and a few GAN-based options [35].

Generate data transformations. Ratner et al. learned generative sequence model over user-defined transformations using GAN-like framework. Their idea is to compose and parameterize a set of user-specified transformation functions in ways that are diverse but still preserve class labels. Their approach allows leveraging domain knowledge flexibly and straightforwardly. [27].

Cubuk et al. proposed a new procedure called AutoAugment[6] to learn augmentation policies that lead to the highest accuracy of the image classification model on a given dataset. They created a search space of data augmentation policies and used a search algorithm based on reinforcement learning to find the optimal one. The results are great: they achieved state-of-the-art accuracy on CIFAR-10, CIFAR-100, SVHN, and ImageNet datasets. Moreover, It is shown that policies learned from one dataset can be transferred to other similar datasets. One of the drawbacks of AutoAugment is the computational complexity and long training time due to the extensive search space of possible policies [6]. Many interesting works have been published trying to address this issue, including Fast AutoAugment [22] and PBA (Population Based Augmentation) [13], which achieved similar results to AutoAugment but using much more efficient algorithms.

3 Proposed Method

We present an intuitive approach for learning data transformations that can be safely used during the model training. Our learned set of augmentations is called safe augmentations, which can be used either for model fine-tuning or along with other augmentation techniques.

Our method does not require substantial computational resources and can be easily performed along with the main task.

3.1 Learning Safe Augmentations

We propose to learn safe augmentations from data using an arbitrary convolutional neural network (CNN). Consider a dataset $D$ and a set of all available augmentation techniques $A$. The task is to define which transformations from set $A$ do not change the distribution of the $D$, i.e., to select $S\subset A$, where $S$ is a set of safe augmentation. Our pipeline can be divided into four main steps.

• •

Step 1. Train the CNN to solve the following multi-label classification problem. Given a set $A$, for every batch of images, a random subset $a\subset A$ is applied. The subset $a$ is of random size from [math] to the defined maximal size. In our experiments, we used maximum size of 5. Each transformation from subset $a$ is applied with the probability $p=1$. The model tries to predict which augmentations were applied. As a loss function, we use $L_{augm}$ - a multi-label one-versus-all loss based on max-entropy, between input $x$ and target $y$. $L_{augm}$ is equivalent to applying $sigmoid$ function along with the binary cross-entropy loss.

• •

Step 2. After the model is being trained, evaluate it on the unseen test data without any augmentations and collect per-label false positives, i.e., how many times the model predicts that the specific augmentation technique was applied when, in fact, it was not.

• •

Step 3. Evaluate the model on the unseen test data using the same procedure of applying augmentations as in the training phase. Collect per-label classification accuracy for each transformation technique.

• •

Step 4. Divide all augmentations into two groups: safe and others. If the model fails to distinguish whether a particular transformation was applied and the transformation is never predicted on the clean set, then this transformation does not change the data distribution and can be safely used during the training of the original task. Thus, we consider augmentation as safe if it has relatively low per-label classification accuracy on the test set with augmentations (step 3) and low false positive rate on the clean test set without augmentations (step 2).

Example. Figure 2 shows the described above metrics for the Tiny ImageNet dataset along with the image classification accuracy of every single augmentation. $Blur$ is an example of non-safe transformation with the low false positive rate on the clean test set and very high augmentation classification accuracy on the augmented test set, so the model can predict when $Blur$ was applied. On the other hand, $HorizontalFlip$ is an example of safe augmentation with both low false positive rate and augmentation classification accuracy. It is clearly shown than $HorizontalFlip$ significantly increases the image classification accuracy on the Tiny ImageNet dataset, whereas $Blur$ decreases it.

Note that we can only evaluate one transformation at a time, i.e., we cannot take into account the impact of different augmentations on each other. So the combination of safe augmentations is not necessarily safe. For example, given a dataset of $32x32$ images and image classification task, our method found that RandomCrop(25, 25) and CenterCrop(25, 25) are safe augmentations. However, when these two functions are applied together, it is likely that such a combination is no longer safe because the augmented image could be too small.

3.2 Joint Learning

To learn augmentations that not only do not change the data distribution but also improve original task accuracy(e.g., image classification, image segmentation), we trained the multi-label classification problem (step 1) in a joint learning setup. To do that, we propose to modify the architecture of the original models in the following way.

For the image classification task, the new loss $L_{total}$ is calculated as sum of the augmentation classification loss $L_{augm}$ and the image classification loss $L_{class}$.

[TABLE]

where $L_{class}$ is the cross-entropy loss.

For the semantic image segmentation task, the new loss $L_{total}$ is calculated in a similar way as a sum of the augmentation classification loss $L_{augm}$ and the semantic segmentation loss $L_{segm}$.

[TABLE]

where $L_{segm}$ is the cross-entropy loss same as $L_{class}$, but here $x$ is a two-dimensional predicted mask, $y$ is the two-dimensional target mask, and the goal is to label every pixel in $x$ with the correct class.

For each augmented batch of images, we calculate defined above $L_{total}$ loss and then perform gradient updates. Joint learning setup helps find a better set of safe augmentations that can be used to improve the performance of the original task. All the results presented were obtained using this approach.

3.3 Using Safe Augmentations

Having learned the set of safe augmentations $S$ for a given dataset and task, we propose to use them in the following two ways:

• •

For model fine-tuning

Step 1. Train the original task using a set of all augmentations $A$. For every batch of images, a random subset $a\subset A$ of fixed size is applied. Each transformation is applied with the fixed probability $p=0.5$.

Step 2. Fine-tune the already pre-trained model on all augmentations using a subset of safe augmentations $S$. For every batch of images, a random subset $s\subset S$ of fixed size is applied. The subset size and probability of applying transformations are the same as in the previous step.

We believe that using all augmentations, including those that change the data distribution can force the model to learn more general features. Thus, they can be used for the model pre-training. After that, we need to fine-tune the model using safe augmentations for learning dataset-specific features.

• •

Along with other augmentation techniques

Safe augmentations alone cannot provide enough generalization. In our experiments, we show that using safe augmentations along with other augmentation techniques, e.g., baseline augmentations and Cutout [9] leads to better results.

4 Experimental Evaluation

4.1 Implementation Details

We implemented all of our models using PyTorch [1] [25]. All models were trained on a single GTX 1080 GPU. We perform the phase of learning safe augmentations using the joint learning setup.

Augmentations. We use a main set $A$ of $15$ common augmentations (see Table 1) from albumentations library that provides fast image transforms based on highly-optimized OpenCV library [4].

Image classification. As the main model, we use DenseNet-121 [15]. Both learning safe augmentations and image classification tasks are trained from scratch using stochastic gradient descent (SGD) optimizer using the batch of size 256. The initial learning rate is set to $10^{-1}$, momentum to $0.9$ and weight decay to $0.0005$. All models are trained for $500$ epochs with reducing learning rate on the plateau by $0.1$ with $10$ epochs patience and early stopping with $20$ epochs patience.

Image segmentation. As the main model, we use Feature Pyramid Network (FPN) [23] with the DenseNet-121 [15] backbone. Both learning safe augmentations and image classification tasks are trained from scratch using Adam [16] optimizer. The initial learning rate is set to $10^{-4}$. All models are trained for $200$ epochs with reducing learning rate on plateau by $0.5$ with $7$ epochs patience and early stopping with $15$ epochs patience.

4.2 Augmentation Subset Size

We investigate impact of different subset sizes on image classification accuracy (see Figure 3). Empirically proven that subset of size 3 leads to the best image classification accuracy for both all augmentations set $A$ and safe augmentations set $S$. Hence, in all our experiments we are using $a\subset A$ and $s\subset S$ of size 3.

4.3 Quantitative Evaluation on Image Classification

We evaluate our method on 4 popular image classification datasets, namely CIFAR-10 [17], CIFAR-100 [18], SVHN [24] and Tiny ImageNet. For every dataset, we train the image classification task using different augmentation techniques and models. All augmentations in our image classification experiments are applied with the fixed probability $p=0.5$ and with default magnitude.

• •

CIFAR-10 and CIFAR-100. We used all crops described in the Table 1 with the size $(25x25)$. As a baseline augmentations, we used horizontal flips with 50 % probability, zero-padding and random crops, which are conventional for these datasets [11]. The training data is also normalized by the respective dataset statistics.

Our method found almost the same set of safe augmentations for CIFAR-10 and CIFAR-100 datasets, which makes sense because these datasets are very similar (see Table 1). We manually defined another set called Safe v2 by removing RandomCrop and CenterCrop, since our approach can only evaluate transformations independently.

Fine-tuning using set Safe v2 achieves the best accuracy both on CIFAR-10 and CIFAR-100 comparing to other evaluated augmentation approaches, namely without augmentations, baseline, only safe, only all and different fine-tuned setups (see Table 2).

• •

SVHN. Same as for CIFAR-10 and CIFAR-100 we used all crops described in the Table 1 with the size $(25x25)$. As a baseline augmentations, we used zero-padding and random crops. The training data is normalized by the respective dataset statistics.

Our method found four safe augmentations (see Table 1), which all are color-based transformations and can be safely used with digits.

Fine-tuning using Safe set produces the best accuracy both for DenseNet-121 and DenseNet-169 models comparing to other evaluated augmentation approaches (see Table 2).

• •

Tiny ImageNet. We used all crops described in the Table 1 with the size $50x50$. As a baseline augmentations, we used horizontal flips with $p=0.5$ and random distortions of colors, which are a standard data augmentation techniques for the ImageNet dataset [7, 14].

Our method found the same safe augmentations (see Table 1) as for CIFAR-10. We manually defined another set called Safe v2 in the same way as for CIFAR-10 and CIFAR-100.

Fine-tuning using Safe set again achieves the best accuracy comparing to other evaluated augmentation approaches (see Table 2).

Safe Augmentation vs AutoAugment. We also evaluated our second proposed way of using safe augmentations along with other augmentation techniques. We compare our approach with AutoAugment [6] and Cutout [9] on CIFAR-10, CIFAR-100 and SVHN datasets. For CIFAR-10 and CIFAR-100 we use cutout of size $16x16$ and for SVHN of size $20x20$ the same way as AutoAugment does. Safe augmentations are applied along with the baseline and Cutout.

The results on Table 5 show that using safe augmentations in both fine-tuning way and along with other augmentation techniques lead to better results.

4.4 Quantitative Evaluation on Image Segmentation

We also evaluate our method on Cityscapes - a popular image segmentation dataset [5]. We train both augmentation classification task and image classification task using the batch of size $16$. We rescale every image to $256x256$ pixels due to the limited training resources.

We also change the size of the crops to $512x512$. All augmentations in our experiments with the image classification task are applied with the fixed probability $p=0.5$ and with default magnitude.

• •

CityScapes. As a baseline augmentations, we used horizontal flips with $p=0.5$ and rotation with $p=0.5$ with angle chosen randomly from 0 to 20 degrees, Our method found four safe augmentations (see Table 1).

Fine-tuning using Safe set again achieves the best accuracy comparing to other evaluated augmentation approaches (see Table 4).

5 Conclusion

This paper introduces Safe Augmentation, a simple yet efficient algorithm for automatic selection of data augmentations with promising results on quantitative benchmarks. In addition to the simplicity, explainability and flexibility S̱afe Augmentation also introduce two different ways of using learned augmentations along with other augmentations techniques. We plan to extend Safe Augmentation further for automatic parameter selection, as well as for adversarial robustness.

6 Acknowledgement

The authors were supported by the WANNABY, SoftServe and Ukrainian Catholic University. We thank ods.ai community for the constructive feedback, advice and ideas.