{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "UrQeEKgcumro"
},
"source": [
"# MixUp augmentation for image classification\n",
"\n",
"**Author:** [Sayak Paul](https://twitter.com/RisingSayak)
\n",
"**Date created:** 2021/03/06
\n",
"**Last modified:** 2023/07/24
\n",
"**Description:** Data augmentation using the mixup technique for image classification."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RX6DeC3Cumrp"
},
"source": [
"## Introduction"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "J7gTGXcXumrq"
},
"source": [
"_mixup_ is a *domain-agnostic* data augmentation technique proposed in [mixup: Beyond Empirical Risk Minimization](https://arxiv.org/abs/1710.09412)\n",
"by Zhang et al. It's implemented with the following formulas:\n",
"\n",
"\n",
"\n",
"(Note that the lambda values are values with the [0, 1] range and are sampled from the\n",
"[Beta distribution](https://en.wikipedia.org/wiki/Beta_distribution).)\n",
"\n",
"The technique is quite systematically named. We are literally mixing up the features and\n",
"their corresponding labels. Implementation-wise it's simple. Neural networks are prone\n",
"to [memorizing corrupt labels](https://arxiv.org/abs/1611.03530). mixup relaxes this by\n",
"combining different features with one another (same happens for the labels too) so that\n",
"a network does not get overconfident about the relationship between the features and\n",
"their labels.\n",
"\n",
"mixup is specifically useful when we are not sure about selecting a set of augmentation\n",
"transforms for a given dataset, medical imaging datasets, for example. mixup can be\n",
"extended to a variety of data modalities such as computer vision, naturallanguage\n",
"processing, speech, and so on."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fxAq9sBlumrr"
},
"source": [
"## Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ukjDkEmyumrs"
},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
"\n",
"import numpy as np\n",
"import keras\n",
"import matplotlib.pyplot as plt\n",
"\n",
"from keras import layers\n",
"\n",
"# TF imports related to tf.data preprocessing\n",
"from tensorflow import data as tf_data\n",
"from tensorflow import image as tf_image\n",
"from tensorflow.random import gamma as tf_random_gamma\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TleJcTlmumrt"
},
"source": [
"## Prepare the dataset\n",
"\n",
"In this example, we will be using the [FashionMNIST](https://github.com/zalandoresearch/fashion-mnist) dataset. But this same recipe can\n",
"be used for other classification datasets as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4bfDyhtfumru"
},
"outputs": [],
"source": [
"(x_train, y_train), (x_test, y_test) = keras.datasets.fashion_mnist.load_data()\n",
"\n",
"x_train = x_train.astype(\"float32\") / 255.0\n",
"x_train = np.reshape(x_train, (-1, 28, 28, 1))\n",
"y_train = keras.ops.one_hot(y_train, 10)\n",
"\n",
"x_test = x_test.astype(\"float32\") / 255.0\n",
"x_test = np.reshape(x_test, (-1, 28, 28, 1))\n",
"y_test = keras.ops.one_hot(y_test, 10)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GG2Kafsrumrv"
},
"source": [
"## Define hyperparameters"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "__uledHZumrv"
},
"outputs": [],
"source": [
"AUTO = tf_data.AUTOTUNE\n",
"BATCH_SIZE = 64\n",
"EPOCHS = 10"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kI2M5KLwumry"
},
"source": [
"## Convert the data into TensorFlow `Dataset` objects"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rE8Yv_l6umr0"
},
"outputs": [],
"source": [
"# Put aside a few samples to create our validation set\n",
"val_samples = 2000\n",
"x_val, y_val = x_train[:val_samples], y_train[:val_samples]\n",
"new_x_train, new_y_train = x_train[val_samples:], y_train[val_samples:]\n",
"\n",
"train_ds_one = (\n",
" tf_data.Dataset.from_tensor_slices((new_x_train, new_y_train))\n",
" .shuffle(BATCH_SIZE * 100)\n",
" .batch(BATCH_SIZE)\n",
")\n",
"train_ds_two = (\n",
" tf_data.Dataset.from_tensor_slices((new_x_train, new_y_train))\n",
" .shuffle(BATCH_SIZE * 100)\n",
" .batch(BATCH_SIZE)\n",
")\n",
"# Because we will be mixing up the images and their corresponding labels, we will be\n",
"# combining two shuffled datasets from the same training data.\n",
"train_ds = tf_data.Dataset.zip((train_ds_one, train_ds_two))\n",
"\n",
"val_ds = tf_data.Dataset.from_tensor_slices((x_val, y_val)).batch(BATCH_SIZE)\n",
"\n",
"test_ds = tf_data.Dataset.from_tensor_slices((x_test, y_test)).batch(BATCH_SIZE)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sJgLEhULumr1"
},
"source": [
"## Define the mixup technique function\n",
"\n",
"To perform the mixup routine, we create new virtual datasets using the training data from\n",
"the same dataset, and apply a lambda value within the [0, 1] range sampled from a [Beta distribution](https://en.wikipedia.org/wiki/Beta_distribution)\n",
"— such that, for example, `new_x = lambda * x1 + (1 - lambda) * x2` (where\n",
"`x1` and `x2` are images) and the same equation is applied to the labels as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "B2T0airYumr3"
},
"outputs": [],
"source": [
"\n",
"def sample_beta_distribution(size, concentration_0=0.2, concentration_1=0.2):\n",
" gamma_1_sample = tf_random_gamma(shape=[size], alpha=concentration_1)\n",
" gamma_2_sample = tf_random_gamma(shape=[size], alpha=concentration_0)\n",
" return gamma_1_sample / (gamma_1_sample + gamma_2_sample)\n",
"\n",
"\n",
"def mix_up(ds_one, ds_two, alpha=0.2):\n",
" # Unpack two datasets\n",
" images_one, labels_one = ds_one\n",
" images_two, labels_two = ds_two\n",
" batch_size = keras.ops.shape(images_one)[0]\n",
"\n",
" # Sample lambda and reshape it to do the mixup\n",
" l = sample_beta_distribution(batch_size, alpha, alpha)\n",
" x_l = keras.ops.reshape(l, (batch_size, 1, 1, 1))\n",
" y_l = keras.ops.reshape(l, (batch_size, 1))\n",
"\n",
" # Perform mixup on both images and labels by combining a pair of images/labels\n",
" # (one from each dataset) into one image/label\n",
" images = images_one * x_l + images_two * (1 - x_l)\n",
" labels = labels_one * y_l + labels_two * (1 - y_l)\n",
" return (images, labels)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LQq9XFpHumr5"
},
"source": [
"**Note** that here , we are combining two images to create a single one. Theoretically,\n",
"we can combine as many we want but that comes at an increased computation cost. In\n",
"certain cases, it may not help improve the performance as well."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EbLOG2vcumr5"
},
"source": [
"## Visualize the new augmented dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8Aex9C0xumr6"
},
"outputs": [],
"source": [
"# First create the new dataset using our `mix_up` utility\n",
"train_ds_mu = train_ds.map(\n",
" lambda ds_one, ds_two: mix_up(ds_one, ds_two, alpha=0.2),\n",
" num_parallel_calls=AUTO,\n",
")\n",
"\n",
"# Let's preview 9 samples from the dataset\n",
"sample_images, sample_labels = next(iter(train_ds_mu))\n",
"plt.figure(figsize=(10, 10))\n",
"for i, (image, label) in enumerate(zip(sample_images[:9], sample_labels[:9])):\n",
" ax = plt.subplot(3, 3, i + 1)\n",
" plt.imshow(image.numpy().squeeze())\n",
" print(label.numpy().tolist())\n",
" plt.axis(\"off\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DILI0-5mumr7"
},
"source": [
"## Model building"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Bn1FoKl5umr8"
},
"outputs": [],
"source": [
"\n",
"def get_training_model():\n",
" model = keras.Sequential(\n",
" [\n",
" layers.Input(shape=(28, 28, 1)),\n",
" layers.Conv2D(16, (5, 5), activation=\"relu\"),\n",
" layers.MaxPooling2D(pool_size=(2, 2)),\n",
" layers.Conv2D(32, (5, 5), activation=\"relu\"),\n",
" layers.MaxPooling2D(pool_size=(2, 2)),\n",
" layers.Dropout(0.2),\n",
" layers.GlobalAveragePooling2D(),\n",
" layers.Dense(128, activation=\"relu\"),\n",
" layers.Dense(10, activation=\"softmax\"),\n",
" ]\n",
" )\n",
" return model\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Pnb-rIgbumr9"
},
"source": [
"For the sake of reproducibility, we serialize the initial random weights of our shallow\n",
"network."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "dNMMCzA2umr-"
},
"outputs": [],
"source": [
"initial_model = get_training_model()\n",
"initial_model.save_weights(\"initial_weights.weights.h5\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-VaVA8cMumr-"
},
"source": [
"## 1. Train the model with the mixed up dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yxs1ZNOPumr-"
},
"outputs": [],
"source": [
"model = get_training_model()\n",
"model.load_weights(\"initial_weights.weights.h5\")\n",
"model.compile(loss=\"categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
"model.fit(train_ds_mu, validation_data=val_ds, epochs=EPOCHS)\n",
"_, test_acc = model.evaluate(test_ds)\n",
"print(\"Test accuracy: {:.2f}%\".format(test_acc * 100))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zd7zeevLumr_"
},
"source": [
"## 2. Train the model *without* the mixed up dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "LAp5QnaKumsA"
},
"outputs": [],
"source": [
"model = get_training_model()\n",
"model.load_weights(\"initial_weights.weights.h5\")\n",
"model.compile(loss=\"categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
"# Notice that we are NOT using the mixed up dataset here\n",
"model.fit(train_ds_one, validation_data=val_ds, epochs=EPOCHS)\n",
"_, test_acc = model.evaluate(test_ds)\n",
"print(\"Test accuracy: {:.2f}%\".format(test_acc * 100))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3sjtf9f8umsN"
},
"source": [
"Readers are encouraged to try out mixup on different datasets from different domains and\n",
"experiment with the lambda parameter. You are strongly advised to check out the\n",
"[original paper](https://arxiv.org/abs/1710.09412) as well - the authors present several ablation studies on mixup\n",
"showing how it can improve generalization, as well as show their results of combining\n",
"more than two images to create a single one."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8i_PuC8bumsP"
},
"source": [
"## Notes\n",
"\n",
"* With mixup, you can create synthetic examples — especially when you lack a large\n",
"dataset - without incurring high computational costs.\n",
"* [Label smoothing](https://www.pyimagesearch.com/2019/12/30/label-smoothing-with-keras-tensorflow-and-deep-learning/) and mixup usually do not work well together because label smoothing\n",
"already modifies the hard labels by some factor.\n",
"* mixup does not work well when you are using [Supervised Contrastive\n",
"Learning](https://arxiv.org/abs/2004.11362) (SCL) since SCL expects the true labels\n",
"during its pre-training phase.\n",
"* A few other benefits of mixup include (as described in the [paper](https://arxiv.org/abs/1710.09412)) robustness to\n",
"adversarial examples and stabilized GAN (Generative Adversarial Networks) training.\n",
"* There are a number of data augmentation techniques that extend mixup such as\n",
"[CutMix](https://arxiv.org/abs/1905.04899) and [AugMix](https://arxiv.org/abs/1912.02781)."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"name": "mixup",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.0"
}
},
"nbformat": 4,
"nbformat_minor": 0
}