Augmentation Methods Using Albumentations And PyTorch | by Varun Dutt | Towards Data Science

Ever since the Deep Neural Net's rise to fame in the late 1990s, limited data has been a stumbling block. Lack of data results in overfitting, especially for architectures with a lot of parameters. Luckily, augmentation has turned out to be a revelation for machine learning tasks with limited data. It offers two major advantages.
First, you can increase the amount of data by creating several augmented copies of your original dataset which can also be used to balance a skewed dataset.
Second, it forces the model to be invariant to non-relevant features in the samples, for example, background in a face detection task. This helps the model generalize better.
In this post, we will explore the latest data augmentation methods and a novel implementation using the methods discussed. We’ll mainly be covering two methods, AutoAugment, and RandAugment. So, let’s start with a brief introduction to both the methods and then move on to the implementation.
AutoAugment introduced in “AutoAugment: Learning Augmentation Strategies from Data” attempts to automate choosing the type and magnitude of the transformations to apply to a sample. They propose a reinforcement learning method that finds an effective data augmentation policy by discretizing the search problem. The proposed method finds an augmentation policy S which has information on the transformations, the magnitudes of the transformations, and the probabilities of using those transformations.
Let us look at some of the details of the algorithm.
A policy S contains five sub-policies, each sub_policy has two transformations with each transformation having a magnitude and a probability parameter.
Note that the sub-policies here are sensitive to the sequence of the transformations and treated as a different sub-policy if the sequence of application is reversed.
There are 16 augmentation transformations that the controller can choose from (will be discussed in the implementation section). The range of the magnitude and the probability of application for each transformation has been discretized into 10 and 11 uniformly spaced values respectively.
The search algorithm has two components: a controller, which is a recurrent neural network(RNN), and the training algorithm, which is the Proximal Policy Optimization algorithm.
The controller (RNN) predicts an augmentation policy S from the search space, the predicted policy is then used to train a much smaller child model. The child model is trained on augmented data generated by applying the 5 predicted sub-policies on the training set. For each example in the mini-batch, one of the 5 sub-policies is chosen randomly to augment the image. The child model is then evaluated on a separate held-out validation set to measure the accuracy R, which is used as the reward signal to train the RNN controller.
You can read the entire paper here for all the details.
AutoAugment was one of the first papers to implement an automated approach to data augmentation effectively and showed substantially better results than previous automated approaches. However, the separate optimization procedure and the prohibitively large search space have a high computational cost. Moreover, the method makes a strong assumption that the result from a smaller network is transferable to a much larger network. This assumption, as it turns out, is not entirely true as pointed out in the RandAugment paper which proposes a much faster and in some situations better-performing method for data augmentation which we will discuss next.
Most learned augmentation algorithms like AutoAugment creates a proxy task that optimizes for the best augmentation policy on a smaller network, as discussed in the previous section this approach has some fundamental flaws which need to be addressed. RandAugment introduced in “RandAugment: Practical automated data augmentation with a reduced search space” attempts to circumvent this proxy task and also drastically reduce the search space to ease the computational load on the training procedure. The method reduces the number of parameters to two, N and M (which will be discussed in detail in the subsequent sub-section), and treats them as hyperparameters in the training procedure.
RandAugment instead of learning a policy picks a transformation randomly from a set of K transformations with a uniform probability of 1/K. The number of transformations to be picked is given by N which is a hyperparameter, this reduces the total policies in the search space to K pow N. The main benefit of the learned augmentation approach comes from increasing the diversity of the samples, uniform probability approach employed by RandAugment maintains the diversity of the samples but with significantly reduced search space.
Finally, the range of magnitude for each transformation similar to AutoAugment is discretized into 10 uniformly spaced values given by M, however, contrary to the learned approach all the transformations follow the same schedule for M which reduces the search space for magnitudes to just one. This is due to the fact that the optimal magnitude for augmentation depends on the size of the model and the training set, larger networks require larger data distortions(please refer to the image above) which intuitively makes sense. The paper also points out that larger datasets require larger data distortions which seems odd, one hypothesis that the authors put forward is that aggressive augmentation in small datasets might lead to a low signal-to-noise ratio. In any case, there is a clear pattern for optimal magnitude that all transformations follow based on the model size and training set size, this exposes a fundamental flaw in the learned augmentation approach which uses smaller networks to estimate transformation magnitude for all the tasks.
RandAugment uses a simple grid search to find optimal values for N and M, this is done by selecting a set of values for both the hyperparameters and testing all the permutations from those sets on a separate validation set. Finally, choosing the combination that gives the best improvement on the validation set. This procedure can be performed for every sample, every batch, or every epoch.
Now, that we have a fairly good understanding of the algorithms let us move towards the implementation. As discussed previously RandAugment is a much better algorithm, however, some key findings from the AutoAugment paper can be used in conjunction with RandAugment to reduce the computational cost even further.
You can download the notebook for this post from here, the notebook has all the necessary code to set up the augmentation algorithm and test it with Resnet50 on a small part of the Imagenet dataset provided by fastai. However, we will be only focusing on the part of the notebook relevant to augmentation in this post.
Let’s start by looking at the custom RandAugment function from the notebook and then break it down.
def randAugment(N, M, p, mode="all", cut_out = False):  # Magnitude(M) search space    shift_x = np.linspace(0,150,10)  shift_y = np.linspace(0,150,10)  rot = np.linspace(0,30,10)  shear = np.linspace(0,10,10)  sola = np.linspace(0,256,10)  post = [4,4,5,5,6,6,7,7,8,8]  cont = [np.linspace(-0.8,-0.1,10),np.linspace(0.1,2,10)]  bright = np.linspace(0.1,0.7,10)  shar = np.linspace(0.1,0.9,10)  cut = np.linspace(0,60,10)  # Transformation search space  Aug =[#0 - geometrical        A.ShiftScaleRotate(shift_limit_x=shift_x[M], rotate_limit=0,   shift_limit_y=0, shift_limit=shift_x[M], p=p),        A.ShiftScaleRotate(shift_limit_y=shift_y[M], rotate_limit=0, shift_limit_x=0, shift_limit=shift_y[M], p=p),        A.IAAAffine(rotate=rot[M], p=p),        A.IAAAffine(shear=shear[M], p=p),        A.InvertImg(p=p),        #5 - Color Based        A.Equalize(p=p),        A.Solarize(threshold=sola[M], p=p),        A.Posterize(num_bits=post[M], p=p),        A.RandomContrast(limit=[cont[0][M], cont[1][M]], p=p),        A.RandomBrightness(limit=bright[M], p=p),        A.IAASharpen(alpha=shar[M], lightness=shar[M], p=p)]  # Sampling from the Transformation search space  if mode == "geo":     ops = np.random.choice(Aug[0:5], N)  elif mode == "color":     ops = np.random.choice(Aug[5:], N)  else:    ops = np.random.choice(Aug, N)    if cut_out:    ops.append(A.Cutout(num_holes=8, max_h_size=int(cut[M]),   max_w_size=int(cut[M]), p=p))  transforms = A.Compose(ops)  return transforms, ops
The function can roughly be divided into the following two sections.
This section discretizes the magnitude range into 10 uniformly spaced values.
“Aug” is the Transformation Search Space from which the algorithm can pick N transformations with uniform probability. Please take the time to see what these individual transformations do to get a better feel for creating this search space.
Note that the function has three modes “geo”, “color” and “all”. The “geo” mode is for tasks where color-based features define the object of interest in the image eg. fire detection. Similarly, the “color” mode is for tasks where the shape defines the object of interest eg. face detection, detecting a car, etc. Whereas, the “all” mode uses all the transformations.
The authors in the AutoAugment paper showed that different datasets were sensitive to different transformations in both positive and negative directions. Therefore, the “No Free Lunch” theorem applies in data augmentation as well and you will need to tweak the algorithm for every project.
Note that the magnitude ranges given in the table above are for functions from the PIL library, however, we are using Albumentaions in the randAugment() function, mainly because Albumentaions library uses open cv which is faster compared to PIL. You would have to experiment and figure out the magnitude range for each transformation for the library you like to use .
Let’s now look at the schedule for the magnitude of the transformations during training.
def train_model(model, criterion, optimizer, scheduler, aug_mode="all", cut_out=False, max_M=9, num_epochs=25):  best_model_wts = copy.deepcopy(model.state_dict())  best_acc = 0.0  N=2;M=0;p=0.5  for epoch in range(num_epochs):        print('Epoch {}/{}'.format(epoch, num_epochs - 1))        print('-' * 10)            transforms, ops = randAugment(N=N, M=M, p=p, mode=aug_mode, cut_out=cut_out)        dataloaders, dataset_sizes = create_dataloaders(root, train_df, valid_df, label_dict, bs=32, transforms=transforms)        # Each epoch has a training and validation phase        for phase in ['train', 'val']:            if phase == 'train':                model.train()  # Set model to training mode            else:                model.eval()   # Set model to evaluate mode            running_loss = 0.0            running_corrects = 0            # Iterate over data.            for inputs, labels in dataloaders[phase]:                inputs = inputs.to(device)                labels = labels.to(device)                # zero the parameter gradients                optimizer.zero_grad()                # forward                # track history if only in train                with torch.set_grad_enabled(phase == 'train'):                    outputs = model(inputs)                    preds = torch.argmax(input = outputs, dim = 1)                    loss = criterion(outputs, labels)                    # backward + optimize only if in training phase                    if phase == 'train':                        loss.backward()                        optimizer.step()                # statistics                running_loss += loss.item() * inputs.size(0)                running_corrects += torch.sum(preds ==          torch.argmax(input = labels, dim = 1))            if phase == 'train':                scheduler.step()            epoch_loss = running_loss / dataset_sizes[phase]            epoch_acc = running_corrects.double() /  dataset_sizes[phase]                print('{} Loss: {:.4f} Acc: {:.4f}'.format(                phase, epoch_loss, epoch_acc))            # deep copy the model            if phase == 'val' and epoch_acc > best_acc:                best_acc = epoch_acc                best_model_wts = copy.deepcopy(model.state_dict())            elif phase == 'val' and epoch_acc < best_acc:                stp = int((9-M)*((best_acc-epoch_acc)/best_acc))                M += max(1,stp)                M = min(M, max_M)                if M < max_M:                  print("Augmentaion Magnitude Changed To : {}\n".format(M))    print('Best val Acc: {:4f}'.format(best_acc))    # load best model weights    model.load_state_dict(best_model_wts)    return model
The lines of code relevant to the schedule have been boldened in the above code block.
I use a predetermined schedule for M in this example, M is incremented whenever the epoch accuracy is lower than the best accuracy on the validation set. The amount that M will be incremented by is calculated by the difference between the maximum possible value(9) and the current value of M scaled by the ratio of the difference between both the accuracies and the best accuracy.
Note that this schedule incurs no extra cost however is very rigid, RandAugment paper uses a grid search for both N and M after every batch to find their optimal values. However, even with the drastic reduction of computaional cost compared to AutoAugment, grid search is still infeasable for someone with a modest setup and a single gpu for most machine learning projects.
The schedule used in this implementation is not the only available schedule, you can experiment with the number of transformations N per sub-policy and also with probabilities for those transformations p. Augmentation is not an exact science, every lock has a different key, you need to find the matching key through experimentation and creative thinking.
AutoAugment Paper: https://arxiv.org/abs/1805.09501
RandAugment Paper: https://arxiv.org/abs/1909.13719