tutorials/image_classification/DISTIL_Example_CIFAR10.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K8AxA7KVfv9m"
      },
      "source": [
        "# **DISTIL Usage Example: CIFAR10**\n",
        "\n",
        "Here, we show how to use DISTIL to perform active learning on image classification tasks (CIFAR10). This notebook can be easily executed on Google Colab."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m4Y2r9Y_fH5B"
      },
      "source": [
        "## Installation and Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9bsVDupO_EOh",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Get DISTIL\n",
        "!git clone https://github.com/decile-team/distil.git\n",
        "!pip install -r distil/requirements/requirements.txt\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "import sys\n",
        "import torch\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from torch import nn\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import Dataset, Subset, ConcatDataset, DataLoader\n",
        "from torchvision import transforms\n",
        "from torchvision.datasets import cifar\n",
        "\n",
        "sys.path.append('distil/')\n",
        "from distil.active_learning_strategies import GLISTER, BADGE, EntropySampling, RandomSampling   # All active learning strategies showcased in this example\n",
        "from distil.utils.models.resnet import ResNet18                                                 # The model used in our image classification example\n",
        "from distil.utils.train_helper import data_train                                                # A utility training class provided by DISTIL\n",
        "from distil.utils.utils import LabeledToUnlabeledDataset                                        # A utility wrapper class that removes labels from labeled PyTorch dataset objects"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vvn_Dal4gb8S"
      },
      "source": [
        "## Preparing CIFAR10\n",
        "\n",
        "The CIFAR10 dataset contains 60,000 32x32 color images in 10 different classes.The 10 different classes represent airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. There are 6,000 images of each class. The training set contains 50,000 images, and the test set contains 10,000 images. Here, we do a simple setup of the CIFAR10 dataset that we will use in this example. More importantly, we define a split on CIFAR10's training set into an initial labeled seed set and an unlabeled set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jdLy9wyjCuT-",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "data_set_name = 'CIFAR10'\n",
        "download_path = '.'\n",
        "\n",
        "# Define transforms on the dataset splits of CIFAR10. Here, we use random crops and horizontal flips for training augmentations.\n",
        "# Both the train and test sets are converted to PyTorch tensors and are normalized around the mean/std of CIFAR-10.\n",
        "cifar_training_transform = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "cifar_test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "\n",
        "# Get the dataset objects from PyTorch. Here, CIFAR10 is downloaded, and the transform is applied when points \n",
        "# are retrieved.\n",
        "cifar10_full_train = cifar.CIFAR10(download_path, train=True, download=True, transform=cifar_training_transform)\n",
        "cifar10_test = cifar.CIFAR10(download_path, train=False, download=True, transform=cifar_test_transform)\n",
        "\n",
        "# Get the dimension of the images. Here, we simply take the very first image of CIFAR10 \n",
        "# and query its dimension.\n",
        "dim = np.shape(cifar10_full_train[0][0])\n",
        "\n",
        "# We now define a train-unlabeled split for the sake of the experiment. Here, we simply take 1000 points as the initial seed set.\n",
        "# The rest of the points are taken as the unlabeled set. While the unlabeled set constructed here technically has labels, they \n",
        "# are only used when querying for labels. Hence, they only exist here for the sake of experimental design.\n",
        "train_size = 1000\n",
        "cifar10_train = Subset(cifar10_full_train, list(range(train_size)))\n",
        "cifar10_unlabeled = Subset(cifar10_full_train, list(range(train_size, len(cifar10_full_train))))\n",
        "\n",
        "# Define the number of active learning rounds to conduct, the budget, and the number of classes in CIFAR10\n",
        "nclasses = 10\n",
        "n_rounds = 9\n",
        "budget = 500 "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JVwJ7RcOhVc7"
      },
      "source": [
        "## Preparing the Model\n",
        "\n",
        "Here, we use DISTIL's provided implementation of the [ResNet-18](https://arxiv.org/abs/1512.03385) architecture. We also create a model directory to store trained models in this example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4pHzwNwvhVts",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "net = ResNet18()\n",
        "base_dir = \"models\"\n",
        "os.makedirs(base_dir, exist_ok = True)\n",
        "model_directory = os.path.join(base_dir, 'base_model.pth')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "spOTFeWgfm4b"
      },
      "source": [
        "## Training an Initial Model\n",
        "Here, we train an initial model. We do so by creating a training loop object on the initial seed set, the model architecture, and a list of provided arguments. We then save the initial model. Note: If you've already run the first cell, then you can simply run the second cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ru5tj6u8fwEi",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Specify additional training parameters. Here, we set the maximum number of epochs of training to 300, \n",
        "# the learning rate to 0.01, the batch size to 20, the maximum train accuracy of training to 0.99, and \n",
        "# the optimizer to stochastic gradient descent.\n",
        "args = {'n_epoch':300, 'lr':float(0.01), 'batch_size':20, 'max_accuracy':0.99, 'optimizer':'sgd'} \n",
        "dt = data_train(cifar10_train, net, args)\n",
        "clf = dt.train()\n",
        "torch.save(clf.state_dict(), model_directory)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dd-ieb7RgyJY",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "base_dir = \"models\"\n",
        "model_directory = os.path.join(base_dir, 'base_model.pth')\n",
        "net.load_state_dict(torch.load(model_directory))\n",
        "clf = net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HhsUx57BjBT2"
      },
      "source": [
        "## Active Learning Strategies\n",
        "\n",
        "Here, we show examples of a couple active learning strategies being used in the setting of image classification."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "W-Bxn6LmCeEI"
      },
      "source": [
        "### Random Sampling\n",
        "This strategy is often used as a baseline, where we pick a subset of unlabeled points randomly. Here, we create a instance of `distil.active_learning_strategies.random_sampling.RandomSampling` by passing following parameters:\n",
        "\n",
        "**training_dataset** – The labeled dataset, which has a PyTorch interface that returns a tensor of data and a tensor of labels OR a dictionary with a \"labels\" key.\n",
        "\n",
        "**unlabeled_dataset** – The unlabeled dataset, which has a PyTorch interface that returns a tensor of data OR a dictionary.\n",
        "\n",
        "**net** – Model architecture used for training. Could be instance of models defined in distil.utils.models or something similar (in this instance, a Huggingface transformer).\n",
        "\n",
        "**nclasses** – No. of classes in the dataset\n",
        "\n",
        "**args**– This dictionary should have ‘batch_size’ as a key. 'batch_size' should be such that one can exploit the benefits of tensorization while honouring the resourse constraits. This ‘batch_size’ therefore can be different than the one used for training.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j8OdeQ63DZuP",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Initialize the random sampling AL strategy. Note: The labels are shaved off the unlabeled dataset above to match the setting.\n",
        "strategy_args = {'batch_size' : 20}\n",
        "strategy = RandomSampling(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled), net, nclasses, strategy_args)\n",
        "\n",
        "# Use the same training parameters as before\n",
        "args = {'n_epoch':300, 'lr':float(0.01), 'batch_size':20, 'max_accuracy':0.99, 'optimizer':'sgd'} \n",
        "dt = data_train(cifar10_train, clf, args)\n",
        "\n",
        "# Update the model used in the AL strategy with the loaded initial model\n",
        "strategy.update_model(clf)\n",
        "\n",
        "# Get the test accuracy of the initial model\n",
        "acc = np.zeros(n_rounds)\n",
        "acc[0] = dt.get_acc_on_set(cifar10_test)\n",
        "print('Initial Testing accuracy:', round(acc[0]*100, 2), flush=True)\n",
        "\n",
        "# User Controlled Loop\n",
        "for rd in range(1, n_rounds):\n",
        "    print('-------------------------------------------------')\n",
        "    print('Round', rd) \n",
        "    print('-------------------------------------------------')\n",
        "\n",
        "    # Use select() to obtain the indices in the unlabeled set that should be labeled\n",
        "    cifar10_full_train.transform = cifar_test_transform       # Disable augmentation while selecting new points as to not interfere with the strategies\n",
        "    idx = strategy.select(budget)\n",
        "    cifar10_full_train.transform = cifar_training_transform   # Enable augmentation\n",
        "\n",
        "    # Add the selected points to the train set. The unlabeled set shown in the next couple lines \n",
        "    # already has the associated labels, so no human labeling is needed. Again, this is because \n",
        "    # we already have the labels a priori. In real scenarios, a human oracle would need to provide \n",
        "    # then before proceeding.\n",
        "    cifar10_train = ConcatDataset([cifar10_train, Subset(cifar10_unlabeled, idx)])\n",
        "    remaining_unlabeled_idx = list(set(range(len(cifar10_unlabeled))) - set(idx))\n",
        "    cifar10_unlabeled = Subset(cifar10_unlabeled, remaining_unlabeled_idx)\n",
        "\n",
        "    print('Number of training points -', len(cifar10_train))\n",
        "\n",
        "    # Update the data used in the AL strategy and the training class\n",
        "    strategy.update_data(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled))\n",
        "    dt.update_data(cifar10_train)\n",
        "\n",
        "    # Retrain the model and update the strategy with the result\n",
        "    clf = dt.train()\n",
        "    strategy.update_model(clf)\n",
        "\n",
        "    # Get new test accuracy\n",
        "    acc[rd] = dt.get_acc_on_set(cifar10_test)\n",
        "    print('Testing accuracy:', round(acc[rd]*100, 2), flush=True)\n",
        "\n",
        "print('Training Completed')\n",
        "\n",
        "# Lastly, we save the accuracies in case a comparison is warranted.\n",
        "with open(os.path.join(base_dir,'random.txt'), 'w') as f:\n",
        "    for item in acc:\n",
        "        f.write(\"%s\\n\" % item)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "jm_ILTinrkLk"
      },
      "source": [
        "### Entropy Sampling\n",
        "A very basic strategy to select unlabeled points is entropy sampling, where we select samples about which the model is most uncertain by measuring the entropy of the class prediction. Hence, a valid strategy is to select those points in the unlabeled set with highest entropy (maximum uncertainty). Specifically, let $z = f(x)$ be the output from the model with class $i$ having score $z_i$. By applying a softmax, we obtain probabilities that we can use: \n",
        "\n",
        "$$\\sigma(z)_i = \\frac{e^{z_i}}{\\sum_j e^{z_j}}$$ \n",
        "\n",
        "Using the softmax probabilities, entropy sampling selects a subset based on the following objective:\n",
        "\n",
        " $$\\text{argmin}_{A \\subseteq U, |A| = k} \\sum_{x \\in A} \\left[ -\\sum_j \\sigma(z)_j*\\log(\\sigma(z)_j) \\right]$$\n",
        "\n",
        "Here, we create a instance of `distil.active_learning_strategies.entropy_sampling.EntropySampling` with the same parameters passed to `distil.active_learning_strategies.random_sampling.RandomSampling`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ieS3rs-TC8bE"
      },
      "source": [
        "**Reloading Base Model & Data**\n",
        "\n",
        "We make sure the fixture is the same by repeating the same setup."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p5sn-ejd2IW4",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "data_set_name = 'CIFAR10'\n",
        "download_path = '.'\n",
        "\n",
        "cifar_training_transform = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "cifar_test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "\n",
        "cifar10_full_train = cifar.CIFAR10(download_path, train=True, download=True, transform=cifar_training_transform)\n",
        "cifar10_test = cifar.CIFAR10(download_path, train=False, download=True, transform=cifar_test_transform)\n",
        "\n",
        "dim = np.shape(cifar10_full_train[0][0])\n",
        "\n",
        "train_size = 1000\n",
        "cifar10_train = Subset(cifar10_full_train, list(range(train_size)))\n",
        "cifar10_unlabeled = Subset(cifar10_full_train, list(range(train_size, len(cifar10_full_train))))\n",
        "\n",
        "nclasses = 10\n",
        "n_rounds = 9    \n",
        "budget = 500 \n",
        "\n",
        "net = ResNet18()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bheCElpz2Q-5",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "base_dir = \"models\"\n",
        "model_directory = os.path.join(base_dir, 'base_model.pth')\n",
        "net.load_state_dict(torch.load(model_directory))\n",
        "clf = net"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_3C0729a2Vzg",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Initialize the random sampling AL strategy. Note: The labels are shaved off the unlabeled dataset above to match the setting.\n",
        "strategy_args = {'batch_size' : 20}\n",
        "strategy = EntropySampling(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled), net, nclasses, strategy_args)\n",
        "\n",
        "# Use the same training parameters as before\n",
        "args = {'n_epoch':300, 'lr':float(0.01), 'batch_size':20, 'max_accuracy':0.99, 'optimizer':'sgd'} \n",
        "dt = data_train(cifar10_train, clf, args)\n",
        "\n",
        "# Update the model used in the AL strategy with the loaded initial model\n",
        "strategy.update_model(clf)\n",
        "\n",
        "# Get the test accuracy of the initial model\n",
        "acc = np.zeros(n_rounds)\n",
        "acc[0] = dt.get_acc_on_set(cifar10_test)\n",
        "print('Initial Testing accuracy:', round(acc[0]*100, 2), flush=True)\n",
        "\n",
        "# User Controlled Loop\n",
        "for rd in range(1, n_rounds):\n",
        "    print('-------------------------------------------------')\n",
        "    print('Round', rd) \n",
        "    print('-------------------------------------------------')\n",
        "\n",
        "    # Use select() to obtain the indices in the unlabeled set that should be labeled\n",
        "    cifar10_full_train.transform = cifar_test_transform       # Disable augmentation while selecting new points as to not interfere with the strategies\n",
        "    idx = strategy.select(budget)\n",
        "    cifar10_full_train.transform = cifar_training_transform   # Enable augmentation\n",
        "\n",
        "    # Add the selected points to the train set. The unlabeled set shown in the next couple lines \n",
        "    # already has the associated labels, so no human labeling is needed. Again, this is because \n",
        "    # we already have the labels a priori. In real scenarios, a human oracle would need to provide \n",
        "    # then before proceeding.\n",
        "    cifar10_train = ConcatDataset([cifar10_train, Subset(cifar10_unlabeled, idx)])\n",
        "    remaining_unlabeled_idx = list(set(range(len(cifar10_unlabeled))) - set(idx))\n",
        "    cifar10_unlabeled = Subset(cifar10_unlabeled, remaining_unlabeled_idx)\n",
        "\n",
        "    print('Number of training points -', len(cifar10_train))\n",
        "\n",
        "    # Update the data used in the AL strategy and the training class\n",
        "    strategy.update_data(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled))\n",
        "    dt.update_data(cifar10_train)\n",
        "\n",
        "    # Retrain the model and update the strategy with the result\n",
        "    clf = dt.train()\n",
        "    strategy.update_model(clf)\n",
        "\n",
        "    # Get new test accuracy\n",
        "    acc[rd] = dt.get_acc_on_set(cifar10_test)\n",
        "    print('Testing accuracy:', round(acc[rd]*100, 2), flush=True)\n",
        "\n",
        "print('Training Completed')\n",
        "\n",
        "#Saving accuracies for further analysis\n",
        "with open(os.path.join(base_dir,'entropy.txt'), 'w') as f:\n",
        "    for item in acc:\n",
        "        f.write(\"%s\\n\" % item)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "BAGqAV0GrwwN"
      },
      "source": [
        "### BADGE\n",
        "This method is based on the paper [Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds](https://arxiv.org/abs/1906.03671). The strategy is meant to select points that are both diverse (e.g., their embeddings span multiple directions) and uncertain (e.g., their contribution to the loss is large). The following steps are taken:\n",
        "\n",
        "* Calculate the pseudo-label for each point in the unlabeled set. The pseudo-label is the class with the highest probability.\n",
        "* Compute the cross-entropy loss for each point in the unlabeled set using this pseudo-label.\n",
        "* Obtain the resulting loss gradients on the last linear layer of the model for each point. (These are referred to as the hypothesized loss gradients.)\n",
        "* Using these gradients as a form of embedding for each unlabeled point, run k-means++ initialization on this embedding set, retrieving $k$ centers. Each center is a point from the unlabeled set, and $k$ represents the active learning budget.\n",
        "* Request labels for the $k$ points whose embeddings were selected.\n",
        "\n",
        "Here we create a instance of `distil.active_learning_strategies.badge.BADGE` with same parameters passed to `distil.active_learning_strategies.random_sampling.RandomSampling`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Z2N1gsCGE8U"
      },
      "source": [
        "**Reloading Base Model & Data**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "voTthJxV2kLB",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "data_set_name = 'CIFAR10'\n",
        "download_path = '.'\n",
        "\n",
        "cifar_training_transform = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "cifar_test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "\n",
        "cifar10_full_train = cifar.CIFAR10(download_path, train=True, download=True, transform=cifar_training_transform)\n",
        "cifar10_test = cifar.CIFAR10(download_path, train=False, download=True, transform=cifar_test_transform)\n",
        "\n",
        "dim = np.shape(cifar10_full_train[0][0])\n",
        "\n",
        "train_size = 1000\n",
        "cifar10_train = Subset(cifar10_full_train, list(range(train_size)))\n",
        "cifar10_unlabeled = Subset(cifar10_full_train, list(range(train_size, len(cifar10_full_train))))\n",
        "\n",
        "nclasses = 10\n",
        "n_rounds = 9    ##Number of rounds to run active learning\n",
        "budget = 500 \n",
        "\n",
        "net = ResNet18()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZF0BAQAL2kLB",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "base_dir = \"models\"\n",
        "model_directory = os.path.join(base_dir, 'base_model.pth')\n",
        "net.load_state_dict(torch.load(model_directory))\n",
        "clf = net"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6x0xEWpl2kLC",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Initialize the random sampling AL strategy. Note: The labels are shaved off the unlabeled dataset above to match the setting.\n",
        "strategy_args = {'batch_size' : 20}\n",
        "strategy = BADGE(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled), net, nclasses, strategy_args)\n",
        "\n",
        "# Use the same training parameters as before\n",
        "args = {'n_epoch':300, 'lr':float(0.01), 'batch_size':20, 'max_accuracy':0.99, 'optimizer':'sgd'} \n",
        "dt = data_train(cifar10_train, clf, args)\n",
        "\n",
        "# Update the model used in the AL strategy with the loaded initial model\n",
        "strategy.update_model(clf)\n",
        "\n",
        "# Get the test accuracy of the initial model\n",
        "acc = np.zeros(n_rounds)\n",
        "acc[0] = dt.get_acc_on_set(cifar10_test)\n",
        "print('Initial Testing accuracy:', round(acc[0]*100, 2), flush=True)\n",
        "\n",
        "# User Controlled Loop\n",
        "for rd in range(1, n_rounds):\n",
        "    print('-------------------------------------------------')\n",
        "    print('Round', rd) \n",
        "    print('-------------------------------------------------')\n",
        "\n",
        "    # Use select() to obtain the indices in the unlabeled set that should be labeled\n",
        "    cifar10_full_train.transform = cifar_test_transform       # Disable augmentation while selecting new points as to not interfere with the strategies\n",
        "    idx = strategy.select(budget)\n",
        "    cifar10_full_train.transform = cifar_training_transform   # Enable augmentation\n",
        "\n",
        "    # Add the selected points to the train set. The unlabeled set shown in the next couple lines \n",
        "    # already has the associated labels, so no human labeling is needed. Again, this is because \n",
        "    # we already have the labels a priori. In real scenarios, a human oracle would need to provide \n",
        "    # then before proceeding.\n",
        "    cifar10_train = ConcatDataset([cifar10_train, Subset(cifar10_unlabeled, idx)])\n",
        "    remaining_unlabeled_idx = list(set(range(len(cifar10_unlabeled))) - set(idx))\n",
        "    cifar10_unlabeled = Subset(cifar10_unlabeled, remaining_unlabeled_idx)\n",
        "\n",
        "    print('Number of training points -', len(cifar10_train))\n",
        "\n",
        "    # Update the data used in the AL strategy and the training class\n",
        "    strategy.update_data(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled))\n",
        "    dt.update_data(cifar10_train)\n",
        "\n",
        "    # Retrain the model and update the strategy with the result\n",
        "    clf = dt.train()\n",
        "    strategy.update_model(clf)\n",
        "\n",
        "    # Get new test accuracy\n",
        "    acc[rd] = dt.get_acc_on_set(cifar10_test)\n",
        "    print('Testing accuracy:', round(acc[rd]*100, 2), flush=True)\n",
        "\n",
        "print('Training Completed')\n",
        "\n",
        "#Saving accuracies for further analysis\n",
        "with open(os.path.join(base_dir,'badge.txt'), 'w') as f:\n",
        "    for item in acc:\n",
        "        f.write(\"%s\\n\" % item)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "DuIa_yJdiDOh"
      },
      "source": [
        "### GLISTER\n",
        "This is implemetation of GLISTER-ACTIVE from the paper [GLISTER: Generalization based Data Subset Selection for Efficient and Robust Learning](https://arxiv.org/abs/2012.10630). `GLISTER` tries to solve the following bi-level optimisation problem:\n",
        "\\begin{equation*}\n",
        "\\overbrace{\\underset{{S \\subseteq {\\mathcal U}, |S| \\leq k}}{\\operatorname{argmax\\hspace{0.7mm}}} LL_V(\\underbrace{\\underset{\\theta}{\\operatorname{argmax\\hspace{0.7mm}}} LL_T( \\theta, S)}_{inner-level}, {\\mathcal V})}^{outer-level}\n",
        "\\end{equation*}\n",
        "where is $S$ is set of points selected at each round, ${\\mathcal V}$ could be a dedicated validation set with labeled points or could be union of labeled and unlabeled points with hypothesized labels, and $k$ is the budget.\n",
        "\n",
        "Solving this problem directly is almost intractable due to the bi-level nature (e.g., repeatedly having to solve the inner problem); therefore, GLISTER resorts to one-step approximations. To start, we set $S^0$ as the empty set and incrementally build it as $S^k = S^{k-1} \\cup e$, where $e$ is $\\underset{e}{\\operatorname{argmax\\hspace{0.7mm}}} G_{\\theta}(e | S^k)$, defined below: $$G_{\\theta}(e | S^k) = LL_{V}(\\theta^{k}, {\\mathcal V})$$\n",
        "\n",
        "Concurrently, GLISTER updates the model parameters via the aforementioned one-step approximation after each greedy selection, which allows for recalculation of $G_{\\theta}$:\n",
        "\n",
        "$$\\theta^k \\leftarrow \\theta^{k-1} + \\eta \\nabla_{\\theta} LL_T(\\hat{\\theta}, e),$$ \n",
        "\n",
        "where $\\hat{\\theta}$ are the parameters of the model at the begining of the selection.\n",
        "\n",
        "To prevent overfitting, we can add regularization to GLISTER, which can be set by **_typeOf_**. **_typeOf_** can be set to **None** for normal GLISTER (default), **'Rand'** for replacing **_lam_** fraction of points with random points, **'Diversity'** for adding a diversity set function while computing gain, and **'FacLoc'** for adding a facility location set function while computing gain. **_lam_** for both **'Diversity'** and **'FacLoc'** determines the weightage given to them while computing the gain.\n",
        "\n",
        "Here, we create a instance of `distil.active_learning_strategies.glister.GLISTER` with the same parameters passed to `distil.active_learning_strategies.random_sampling.RandomSampling`. Additionally, we add to the **_args_** dictionary the required keys ‘batch_size’ and ‘lr’. ‘lr’ should be the learning rate used for training. Lastly, the following parameters can be passed as previously described:\n",
        "\n",
        "**validation_dataset (torch.utils.data.Dataset, optional)** – An optional validation dataset.\n",
        "\n",
        "**typeOf (str, optional)** – Determines the type of regularizer to be used. The default is **None**. For a random regularizer, use **‘Rand’**. To use the facility location set function as a regularizer, use **‘FacLoc’**. To use a diversity set function as a regularizer, use **‘Diversity’**.\n",
        "\n",
        "**lam (float, optional)** – Determines the amount of regularization to be applied. Mandatory if is not **_typeOf_=None** and, by default, is set to **None**. For random regularizers, use values between 0 and 1 as it determines fraction of points replaced by random points. For both **‘Diversity’** and **‘FacLoc’**, **_lam_** determines the weightage given to them while computing the gain.\n",
        "\n",
        "**kernel_batch_size (int, optional)** – For the **'Diversity'** and **'FacLoc'** regularizer versions, a similarity kernel is computed, which entails creating a 3-D Torch tensor of dimensions $kernel\\_batch\\_size^{2}\\times(feature\\ dimension)$. Again, **_kernel_batch_size_** should be such that one can exploit the benefits of tensorization while honouring the resourse constraints."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Hfo3xvHGJog"
      },
      "source": [
        "**Reloading Base Model & Data**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UaVPq19w2sWH",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "data_set_name = 'CIFAR10'\n",
        "download_path = '.'\n",
        "\n",
        "cifar_training_transform = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "cifar_test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])\n",
        "cifar_target_transform = transforms.ToTensor()\n",
        "\n",
        "cifar10_full_train = cifar.CIFAR10(download_path, train=True, download=True, transform=cifar_training_transform, target_transform=torch.tensor)\n",
        "cifar10_test = cifar.CIFAR10(download_path, train=False, download=True, transform=cifar_test_transform, target_transform=torch.tensor)\n",
        "\n",
        "dim = np.shape(cifar10_full_train[0][0])\n",
        "\n",
        "train_size = 1000\n",
        "cifar10_train = Subset(cifar10_full_train, list(range(train_size)))\n",
        "cifar10_unlabeled = Subset(cifar10_full_train, list(range(train_size, len(cifar10_full_train))))\n",
        "\n",
        "nclasses = 10\n",
        "n_rounds = 9    ##Number of rounds to run active learning\n",
        "budget = 500 \n",
        "\n",
        "net = ResNet18()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dbkv74Cq2sWH",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "base_dir = \"models\"\n",
        "model_directory = os.path.join(base_dir, 'base_model.pth')\n",
        "net.load_state_dict(torch.load(model_directory))\n",
        "clf = net"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MFg2EPMC2sWH",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#Initializing Strategy Class\n",
        "strategy_args = {'batch_size' : 20, 'lr' : 0.01}\n",
        "strategy = GLISTER(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled), net, nclasses, strategy_args)\n",
        "\n",
        "# Use the same training parameters as before\n",
        "args = {'n_epoch':300, 'lr':float(0.01), 'batch_size':20, 'max_accuracy':0.99, 'optimizer':'sgd'} \n",
        "dt = data_train(cifar10_train, clf, args)\n",
        "\n",
        "# Update the model used in the AL strategy with the loaded initial model\n",
        "strategy.update_model(clf)\n",
        "\n",
        "# Get the test accuracy of the initial model\n",
        "acc = np.zeros(n_rounds)\n",
        "acc[0] = dt.get_acc_on_set(cifar10_test)\n",
        "print('Initial Testing accuracy:', round(acc[0]*100, 2), flush=True)\n",
        "\n",
        "# User Controlled Loop\n",
        "for rd in range(1, n_rounds):\n",
        "    print('-------------------------------------------------')\n",
        "    print('Round', rd) \n",
        "    print('-------------------------------------------------')\n",
        "\n",
        "    # Use select() to obtain the indices in the unlabeled set that should be labeled\n",
        "    cifar10_full_train.transform = cifar_test_transform       # Disable augmentation while selecting new points as to not interfere with the strategies\n",
        "    idx = strategy.select(budget)\n",
        "    cifar10_full_train.transform = cifar_training_transform   # Enable augmentation\n",
        "\n",
        "    # Add the selected points to the train set. The unlabeled set shown in the next couple lines \n",
        "    # already has the associated labels, so no human labeling is needed. Again, this is because \n",
        "    # we already have the labels a priori. In real scenarios, a human oracle would need to provide \n",
        "    # then before proceeding.\n",
        "    cifar10_train = ConcatDataset([cifar10_train, Subset(cifar10_unlabeled, idx)])\n",
        "    remaining_unlabeled_idx = list(set(range(len(cifar10_unlabeled))) - set(idx))\n",
        "    cifar10_unlabeled = Subset(cifar10_unlabeled, remaining_unlabeled_idx)\n",
        "\n",
        "    print('Number of training points -', len(cifar10_train))\n",
        "\n",
        "    # Update the data used in the AL strategy and the training class\n",
        "    strategy.update_data(cifar10_train, LabeledToUnlabeledDataset(cifar10_unlabeled))\n",
        "    dt.update_data(cifar10_train)\n",
        "\n",
        "    # Retrain the model and update the strategy with the result\n",
        "    clf = dt.train()\n",
        "    strategy.update_model(clf)\n",
        "\n",
        "    # Get new test accuracy\n",
        "    acc[rd] = dt.get_acc_on_set(cifar10_test)\n",
        "    print('Testing accuracy:', round(acc[rd]*100, 2), flush=True)\n",
        "\n",
        "print('Training Completed')\n",
        "\n",
        "#Saving accuracies for further analysis\n",
        "with open(os.path.join(base_dir,'glister.txt'), 'w') as f:\n",
        "    for item in acc:\n",
        "        f.write(\"%s\\n\" % item)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jNZjPPO7java"
      },
      "source": [
        "## Visualizing the Results\n",
        "\n",
        "If all strategies have run to completion, you can run the following cell to view the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6vn65GnhZosp",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Load the accuracies previously obtained\n",
        "with open(os.path.join(base_dir,'entropy.txt'), 'r') as f:\n",
        "  acc_ = f.readlines()\n",
        "acc_en = [round(float(x)*100, 2) for x in acc_]\n",
        "with open(os.path.join(base_dir,'badge.txt'), 'r') as f:\n",
        "  acc_ = f.readlines()\n",
        "acc_bd = [round(float(x)*100, 2) for x in acc_]\n",
        "with open(os.path.join(base_dir,'glister.txt'), 'r') as f:\n",
        "  acc_ = f.readlines()\n",
        "acc_gl = [round(float(x)*100, 2) for x in acc_]\n",
        "with open(os.path.join(base_dir,'random.txt'), 'r') as f:\n",
        "  acc_ = f.readlines()\n",
        "acc_rd = [round(float(x)*100, 2) for x in acc_]\n",
        "\n",
        "# Plot them using matplotlib\n",
        "x_axis = np.array([train_size+budget*i for i in range(n_rounds)])\n",
        "plt.figure()\n",
        "plt.plot(x_axis, acc_gl, 'b-', label='GLISTER RAND',marker='o')\n",
        "plt.plot(x_axis, acc_en, 'g-', label='UNCERTAINITY',marker='o')\n",
        "plt.plot(x_axis, acc_bd, 'c', label='BADGE',marker='o')\n",
        "plt.plot(x_axis, acc_rd, 'r', label='RANDOM',marker='o')\n",
        "\n",
        "plt.legend()\n",
        "plt.xlabel('No of Images')\n",
        "plt.ylabel('Test Accuracy')\n",
        "plt.title('DISTIL_CIFAR10')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "WyIVzbwQgrwN",
        "W-Bxn6LmCeEI",
        "jm_ILTinrkLk",
        "BAGqAV0GrwwN"
      ],
      "name": "DISTIL Example CIFAR10.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}