#!/usr/bin/env python # coding: utf-8 # ## Mixup data augmentation # In[1]: from fastai.gen_doc.nbdoc import * from fastai.callbacks.mixup import * from fastai.vision import * # ## What is mixup? # This module contains the implementation of a data augmentation technique called [mixup](https://arxiv.org/abs/1710.09412). It is extremely efficient at regularizing models in computer vision (we used it to get our time to train CIFAR10 to 94% on one GPU to 6 minutes). # # As the name kind of suggests, the authors of the mixup article propose training the model on mixes of the training set images. For example, suppose we’re training on CIFAR10. Instead of feeding the model the raw images, we take two images (not necessarily from the same class) and make a linear combination of them: in terms of tensors, we have: # # `new_image = t * image1 + (1-t) * image2` # # where t is a float between 0 and 1. The target we assign to that new image is the same combination of the original targets: # # `new_target = t * target1 + (1-t) * target2` # # assuming the targets are one-hot encoded (which isn’t the case in PyTorch usually). And it's as simple as that. # # ![mixup](imgs/mixup.png) # # Dog or cat? The right answer here is 70% dog and 30% cat! # # As the picture above shows, it’s a bit hard for the human eye to make sense of images obtained in this way (although we do see the shapes of a dog and a cat). However, it somehow makes a lot of sense to the model, which trains more efficiently. One important side note is that when training with mixup, the final loss (training or validation) will be higher than when training without it, even when the accuracy is far better: a model trained like this will make predictions that are a bit less confident. # ## Basic Training # To test this method, we first create a [`simple_cnn`](/layers.html#simple_cnn) and train it like we did with [`basic_train`](/basic_train.html#basic_train) so we can compare its results with a network trained with mixup. # In[ ]: path = untar_data(URLs.MNIST_SAMPLE) data = ImageDataBunch.from_folder(path) model = simple_cnn((3,16,16,2)) learn = Learner(data, model, metrics=[accuracy]) # In[ ]: learn.fit(8) # ## Mixup implementation in the library # In the original article, the authors suggest four things: # # 1. Create two separate dataloaders, and draw a batch from each at every iteration to mix them up # 2. Draw a value for t following a beta distribution with a parameter alpha (0.4 is suggested in their article) # 3. Mix up the two batches with the same value t # 4. Use one-hot encoded targets # # This module's implementation is based on these suggestions, and modified where experimental results suggested changes that would improve performance. # The authors suggest using the beta distribution with parameters alpha=0.4. (In general, the beta distribution has two parameters, but in this case they're going to be equal.) Why do they suggest this? Well, with the parameters they suggest, the beta distribution looks like this: # # ![betadist](imgs/betadist-mixup.png) # # meaning that there's a very high probability of picking values close to 0 or 1 (in which case the mixed up image is mostly from only one category) and then a somewhat constant, much smaller probability of picking something in the middle (notice that 0.33 is nearly as likely as 0.5, for instance). # # While this works very well, it’s not the fastest way, and this is the first suggestion we adjust. The unnecessary slowdown with this approach comes from drawing two different batches at every iteration, which means loading twice the number of images and additionally applying any other data augmentation functions to them. To avoid this, we apply mixup on a batch with a shuffled version of itself: this way, the images mixed up are still different. # # Using the same value of `t` for the whole batch is another suggestion we modify. In our experiments, we noticed that the model trained faster if we drew a different `t` for every image in the batch. (Both options got to the same result in terms of accuracy, it’s just that one arrived there more slowly.) # # Finally, notice that with this strategy we might create duplicate images: let’s say we are mixing `image0` with `image1` and `image1` with `image0`, and that we draw `t=0.1` for the first mix and `t=0.9` for the second. Then # # `image0 * 0.1 + shuffle0 * (1-0.1) = image0 * 0.1 + image1 * 0.9` # # and # # `image1 * 0.9 + shuffle1 * (1-0.9) = image1 * 0.9 + image0 * 0.1` # # will be the same. Of course we have to be a bit unlucky for this to happen, but in practice, we saw a drop in accuracy when we didn't remove duplicates. To avoid this, the trick is to replace the vector of `t` we drew with: # # `t = max(t, 1-t)` # # The beta distribution with the two parameters equal is symmetric in any case, and this way we ensure that the largest coefficient is always near the first image (the non-shuffled batch). # ## Adding mixup to the mix # We now add [`MixUpCallback`](/callbacks.mixup.html#MixUpCallback) to our Learner so that it modifies our input and target accordingly. The [`mixup`](/train.html#mixup) function does this for us behind the scenes, along with a few other tweaks described below: # In[ ]: model = simple_cnn((3,16,16,2)) learner = Learner(data, model, metrics=[accuracy]).mixup() learner.fit(8) # Training with mixup improves the best accuracy. Note that the validation loss is higher than without mixup, because the model makes less confident predictions: without mixup, most predictions are very close to 0. or 1. (in terms of probability) whereas the model with mixup makes predictions that are more nuanced. Before using mixup, make sure you know whether it's more important to optimize lower loss or better accuracy. # In[2]: show_doc(MixUpCallback) # Create a [`Callback`](/callback.html#Callback) for mixup on `learn` with a parameter `alpha` for the beta distribution. `stack_x` and `stack_y` determine whether we stack our inputs/targets with the vector lambda drawn or do the linear combination. (In general, we stack the inputs or outputs when they correspond to categories or classes and do the linear combination otherwise.) # ### Callback methods # You don't call these yourself - they're called by fastai's [`Callback`](/callback.html#Callback) system automatically to enable the class's functionality. # In[3]: show_doc(MixUpCallback.on_batch_begin) # Draws a vector of lambda following a beta distribution with `self.alpha` and operates the mixup on `last_input` and `last_target` according to `self.stack_x` and `self.stack_y`. # ## Dealing with the loss # We often have to modify the loss so that it is compatible with mixup. PyTorch was very careful to avoid one-hot encoding targets when possible, so it seems a bit of a drag to undo this. Fortunately for us, if the loss is a classic [cross-entropy](https://pytorch.org/docs/stable/nn.html#torch.nn.functional.cross_entropy), we have # # `loss(output, new_target) = t * loss(output, target1) + (1-t) * loss(output, target2)` # # so we don’t one-hot encode anything and instead just compute those two losses and find the linear combination. # # The following class is used to adapt the loss for mixup. Note that the [`mixup`](/train.html#mixup) function will use it to change the `Learner.loss_func` if necessary. # In[4]: show_doc(MixUpLoss, title_level=3) # ## Undocumented Methods - Methods moved below this line will intentionally be hidden # In[5]: show_doc(MixUpLoss.forward) # In[ ]: