Modern Data Augmentation Techniques for Computer Vision

Theoretical know-how of some modern data augmentations along with there implementations in TensorFlow 2.x.
Some interesting ablation studies.
Comparative study between these techniques.
Benchmarking of models trained with the augmentations techniques on the CIFAR-10-C dataset.

Table of Contents

Experimental Setup and Baseline

Dataset

Our models use the CIFAR-10 dataset for training based on different augmentation strategies.
We use the CIFAR-10-C dataset to measure a model's resilience to data shift. This dataset is constructed by corrupting the original CIFAR test set. Each dataset has 15 noise, blur, weather, and digital corruption types, each appearing at five severity levels. Since these corruptions are used to measure network behavior under data shift, they are not introduced into the training procedure(augmentation). In CIFAR-10-C, the first 10,000 images in each .npy files are the test set images corrupted at severity 1, and the last 10,000 images are the test set images corrupted at severity five.

Architecture

All the experiments use ResNet20 architecture, which uses resnet_v1 model builder from official Keras documentation example on CIFAR-10. You can find the model definition in the repo included in this article.
No data augmentation apart from the augmentation strategy under experimentation was used.
We save the initial weight and train all the models using this initial_wt.h5 to start from the same model weight initialization,
The models were trained with Early Stopping with the patience of 10 epochs and monitored val_loss. The upper bound for epochs was 100.
We save the trained weights and use them for robustness benchmarking. We also save the weights from various ablation studies. You can find all the weights here.

Baseline Model

Cutout Augmentation

Convolutional layers require less regularization since they have much fewer parameters than fully-connected layers.
The second factor is that neighboring pixels in images share much of the same information.

Cutout

Implementation

trainloader = tf.data.Dataset.from_tensor_slices((x_train, y_train))

## Apply cutout augmentation
def augment_cutout(image, label):
  img = tfa.image.random_cutout(image, (10,10), constant_values = 0)
  
  return img, label

## Apply image preprocessing
def preprocess_image(image, label):
  img = tf.cast(image, tf.float32)
  img = img/255.

  return img, label

trainloader = (
    trainloader
    .shuffle(1024)
    .map(preprocess_image, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .map(augment_cutout, num_parallel_calls=AUTO)
    .prefetch(AUTO)
)

I have used tf.data to implement this augmentation technique.
It uses TensorFlow Addon version 0.10.0.

 !pip install tensorflow-addons==0.10.0 

augment_cutout simply wraps tfa.image.random_cutout to apply random cutout to the images. Official documentation is here.
I have used cutsize of (10, 10) as default. You will soon see the dependence of cutsize with model performance.



Comparison with Baseline Model

Unlike the baseline model where overfitting occurred quickly, Cutout augmentation provided necessary regularization enabling longer training time.
Even though the training accuracy with Cutout augmentation is lower than that of the baseline model, the validation accuracy is higher, ensuring better generalization over the training dataset.
The authors of Cutout used standard augmentations like random cropping along with cutout augmentation. On top of that, they used a deeper network and trained with a highly tuned SGD optimizer.

Ablation Study

Cutsize Vs Test Accuracy

ResNet-20 model was trained with 5 different cutsizes.

CUT_SIZES = [4, 8, 12, 16, 20]

They were trained with early-stopping to get the best instance of the model for each cutsize.
The result of this ablation study is shown in the figure below.

Mixup Augmentation

The procedure is dataset-dependent, and thus requires the use of expert knowledge.
Data augmentation assumes that the examples in the vicinity (neighborhood) share the same class, and does not model the vicinity relation across examples of different classes.

mixup



Implementation

def mixup(a, b):
 
  alpha = [0.2]
  beta = [0.2]
 
  # unpack (image, label) pairs
  (image1, label1), (image2, label2) = a, b

  # define beta distribution
  dist = tfd.Beta(alpha, beta)
  # sample from this distribution
  l = dist.sample(1)[0][0]
  
  # mixup augmentation
  img = l*image1+(1-l)*image2
  lab = l*label1+(1-l)*label2

  return img, lab


trainloader1 = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1024).map(preprocess_image, num_parallel_calls=AUTO)
trainloader2 = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1024).map(preprocess_image, num_parallel_calls=AUTO)

trainloader = tf.data.Dataset.zip((trainloader1, trainloader2))
trainloader = (
    trainloader
    .shuffle(1024)
    .map(mixup, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)

I have used tf.data for mixup implementation.
mixup function takes in two (image, label) pairs and applies linear interpolation in image and label space. This function can easily be used for any input modality.
Notice that $λ$ (l) is sampled from a Beta distribution. TensorFlow Probability is used for this.
Since mixup requires two images drawn randomly from the training dataset, trainloader1 and trainloader2 are two separate tf.data datasets zipped using tf.data.Dataset.zip.
This zipped dataset provides two pairs of (image, label) which using map returns augmented image and associated labels.
The implementation is straightforward and has minimal computational overhead.



Comparison with Baseline Model

The baseline model quickly overfitted.
The validation loss followed the training loss without diverging. This allowed for more extended training and far better generalization.
The validation accuracy seems to fluctuate less than the baseline model. mixup helped stabilize the learning process.

So you may ask, what is mixup doing?

The model is encouraged to act linearly in between classes. This linear behavior reduces undesirable oscillations when predicting outside the training examples. This is visible in the validation curves.
mixup provides a smoother(linear) transition from one class decision boundary to another. This provides a better measure of uncertainty.
The model trained with mixup is more stable in terms of model predictions. We will see this soon.

Ablation Study

Alpha Vs Test Accuracy

ResNet-20 was trained with six different alpha values.

ALPHAS = [0.1, 0.2, 0.4, 0.7, 0.9, 1.0]

When alpha=0.0 beta distribution returns 0.5. As alpha increases from 0 towards 1, the shape of the distribution is convex in nature, thus more values lie either towards zero or one. For alpha=1.0, the beta distribution is uniform in nature.
Along the x-axis, x-ticks corresponds to ALPHAS.
The result for this ablation study is shown in the figure below.

CutMix Augmentation

In Cutout augmentation, a region(square) is removed from the image and is filled with either black or grey pixels or with Gaussian noise. This dramatically decreases the portion of informative pixels during the training process. It is a conceptual limitation as CNN is data-hungry.
In Mixup augmentation, we linearly interpolate two randomly drawn images. However, the images generated are somewhat unnatural and confuse the model, especially for the localization task.

CutMix

Instead of removing pixels and filling them with black or grey pixels or Gaussian noise, what if it is replaced with a patch of similar dimensions from another image.
The ground truth labels are mixed proportionally to the number of pixels of combined images. Linearly interpolated label thus produced is proportional to the contributing pixels from the two images.





Implementation

@tf.function
def get_bbox(l):
  cut_rat = tf.math.sqrt(1.-l)

  cut_w = IMG_SHAPE*cut_rat #rw
  cut_w = tf.cast(cut_w, tf.int32)
  
  cut_h = IMG_SHAPE*cut_rat #rh
  cut_h = tf.cast(cut_h, tf.int32)
  
  cx = tf.random.uniform((1,), minval=0, maxval=IMG_SHAPE, dtype=tf.int32) #rx
  cy = tf.random.uniform((1,), minval=0, maxval=IMG_SHAPE, dtype=tf.int32) #ry
  
  bbx1 = tf.clip_by_value(cx[0] - cut_w // 2, 0, IMG_SHAPE)
  bby1 = tf.clip_by_value(cy[0] - cut_h // 2, 0, IMG_SHAPE)
  bbx2 = tf.clip_by_value(cx[0] + cut_w // 2, 0, IMG_SHAPE)
  bby2 = tf.clip_by_value(cy[0] + cut_h // 2, 0, IMG_SHAPE)
  
  target_h = bby2-bby1
  if target_h ==0:
    target_h+=1

  target_w = bbx2-bbx1
  if target_w ==0:
    target_w+=1

  return bbx1, bby1, target_h, target_w

@tf.function
def cutmix(a, b):
  
  (image1, label1), (image2, label2) = a, b

  alpha = [1.]
  beta = [1.]
  
  ## Get sample from beta distribution
  dist = tfd.Beta(alpha, beta)
  ## Lambda
  l = dist.sample(1)[0][0]

  ## Get bbox ofsets and heights and widths
  bbx1, bby1, target_h, target_w = get_bbox(l)

  ## Get patch from image2
  crop2 = tf.image.crop_to_bounding_box(image2, bby1, bbx1, target_h, target_w)
  ## Pad the patch with same offset
  image2 = tf.image.pad_to_bounding_box(crop2, bby1, bbx1, IMG_SHAPE, IMG_SHAPE)
  ## Get patch from image1
  crop1 = tf.image.crop_to_bounding_box(image1, bby1, bbx1, target_h, target_w)
  ## Pad the patch with same offset
  img1 = tf.image.pad_to_bounding_box(crop1, bby1, bbx1, IMG_SHAPE, IMG_SHAPE)

  ## Subtract the patch from image1 so that patch from image2 can be put on instead
  image1 = image1-img1
  ## Add modified image1 and image2 to get cutmix image
  image = image1+image2

  ## Adjust lambda according to pixel ratio
  l = 1 - (target_w * target_h) / (IMG_SHAPE * IMG_SHAPE)
  l = tf.cast(l, tf.float32)

  ## Combine labels
  label = l*label1+(1-l)*label2

  return image, label


trainloader1 = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1024).map(preprocess_image, num_parallel_calls=AUTO)
trainloader2 = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1024).map(preprocess_image, num_parallel_calls=AUTO)

trainloader = tf.data.Dataset.zip((trainloader1, trainloader2))
trainloader = (
    trainloader
    .shuffle(1024)
    .map(cutmix, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)


I have used tf.data to implement this augmentation strategy.
The get_bbox function takes in $λ$(l), which is sampled from a beta distribution and returns bounding box coordinates. It correctly returns (x, y) coordinates for the top-left corner of the patch to be removed along with the height and width of the patch.
cutmix function takes in two (image, label) pairs and applies the CutMix augmentation strategy. First, it samples $λ$(l) from a beta distribution and gets the bounding box coordinates from the get_bbox function. Then we get a crop from image2 and pad it such that the final padded image has this crop at the same spatial location. We also do this for image1, but then we use this to subtract the patch from image1. Finally, we add image1 from which we have removed a patch and image2, which is the patch itself(after cropping+padding operation). Try plotting the generated images for better understanding.
We adjust lambda(l) according to the proportion of the pixels contributed by each image. Moreover, we use this to interpolate the labels linearly.
Since CutMix requires two images drawn randomly from the training dataset, trainloader1 and trainloader2 are two separate tf.data datasets which are zipped using tf.data.Dataset.zip.
This zipped dataset provides two pairs of (image, label) which using map returns an augmented image and associated label.



Comparison with Cutout Augmentation

Cutmix provides stronger regularization compared to Cutout. This can be seen through the loss metrics.
But this strong regularization is indeed helping with reducing generalization error. However, this reduction in validation error is marginal.
For this particular comparison, the model was trained with Cutmix with alpha=1.0. Through the ablation study, we will find out the optimal value of alpha.
The validation metrics are oscillating less when compared to the model trained with cutout augmentation. The oscillation in the case of Cutout can be due to less informative pixels during training, which is not the case with CutMix.
Training with better learning rate scheduling can reduce the risk of overfitting due to strong regularization.

Ablation Study

Alpha VS Test Accuracy

ResNet-20 was trained with six different alpha values.

ALPHAS = [0.1, 0.25, 0.5, 1.0, 2.0, 4.0]

When alpha=0.0 beta distribution returns 0.5. As alpha increases from 0 towards 1, the shape of the distribution is convex in nature, thus more values lie either towards zero or one. For alpha=1.0, the beta distribution is uniform in nature. For alpha>1.0 the shape is concave in nature.
Along the x-axis, x-ticks corresponds to ALPHAS.
The result for this ablation study is shown in the figure.

AugMix Augmentation



AugMix

Augmentations: This method consists of mixing the results from augmentation chains or compositions of augmentation operations. These operations are sampled randomly from a pool of simple augmentation techniques. For operations which can be realized with different level of severity, the level is uniformly sampled. AugMix provides augmentation diversity by having multiple chains of augmentation with varying depths. By default, the number of chains is three, while its depth is uniformly sampled in a range of (1,3). Some augmentation techniques used are:

[autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y, translate_x, translate_y]

Mixing: The images are mixed using the element-wise convex combination for simplicity. After images from different chains are mixed, it is combined with the original image through a second random convex combination sampled from a Beta distribution.
Jensen-Shannon Divergence Consistency Loss: This loss ensures smoother neural network response. Since the semantic content of the image is, to a lot extent, preserved with AugMix, the network should embed $x_{original}$, $x_{augmix1}$, and $x_{augmix2}$ similarly. Thus the original loss is modified by adding this loss. Mathematically,








Implementation



Comparison With Other Augmentation Techniques

By looking at the validation loss for each augmentation strategy, the model performance can be ranked like this.

 augmix > cutmix > mixup > cutout > baseline

The baseline model trained on CIFAR-10 quickly overfitted. Each augmentation strategy is a better addition on top of the baseline model.
CutMix regularized the strongest. We see that mixup was not so aggressive. AugMix initially regularized more strongly but slowly smoothed out. Cutout barely regularized but was a better than baseline.

Model Robustness to CIFAR-10-C

Originally CIFAR-10-C has 19 corruption types. But we will be evaluating on 15 such corruptions as was originally done in the AugMix paper. These corruptions are:

CORRUPTIONS = [
  'gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur',
  'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog',
  'brightness', 'contrast', 'elastic_transform', 'pixelate',
  'jpeg_compression'
]

We are using the best performing models on the CIFAR-10 test dataset trained with each augmentation technique. These models were selected either from the ablation studies(discussed) or by training the model with recommended hyperparameters/settings(AugMix).
For CIFAR-10-C dataset there are five levels of severity for each corruption. We will compute the average error rate for each corruption and then take the average over them.

Summary

Experiment with these augmentation techniques in Colab →