tutorials/image_classification/realistic_scenarios/DISTIL_Example_Rare_Classes_PneumoniaMNIST.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "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.6"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "4e43acf6a8b74698bba7e8ef07039557": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_6cba53ea45c5468fafa46f482f97e517",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_2ca9a395a43e4cc1a7d59a04368c845a",
              "IPY_MODEL_44be0278f4fc47ec8cf438fde4ec12b8",
              "IPY_MODEL_8114758516654ce6972df8a268776717"
            ]
          }
        },
        "6cba53ea45c5468fafa46f482f97e517": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2ca9a395a43e4cc1a7d59a04368c845a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_be7b11692fba4d2593113dfa61c88e4f",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e992b3a0110c4292997f0560fd34f87b"
          }
        },
        "44be0278f4fc47ec8cf438fde4ec12b8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_4be0dd51017b4b879c5b7d8255a88ae5",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 4170669,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 4170669,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_732b2a477e594cac82175561e8574237"
          }
        },
        "8114758516654ce6972df8a268776717": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_02cc627dbedf4361804e840016701708",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 4170752/? [00:01&lt;00:00, 5420110.67it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_97f8e9a087a94a6eb83975bd9c647f78"
          }
        },
        "be7b11692fba4d2593113dfa61c88e4f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "e992b3a0110c4292997f0560fd34f87b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4be0dd51017b4b879c5b7d8255a88ae5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "732b2a477e594cac82175561e8574237": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "02cc627dbedf4361804e840016701708": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "97f8e9a087a94a6eb83975bd9c647f78": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lOkAr6RnQ2BJ"
      },
      "source": [
        "# DISTIL Usage Example: PneumoniaMNIST with Rare Classes\n",
        "\n",
        "Not all datasets have an even spread of classes among the set of labels. Indeed, a dataset might have only a couple elements that have a particular class as a label. Such classes are considered *rare*, and extra work is required to achieve good model performance on these examples. The typical fix is to provide more data with the rare-class label; however, this issue is complicated by the fact that this data is usually part of a massive unlabeled pool of data. Here, we show how to use DISTIL's implementation of [SIMILAR](https://arxiv.org/abs/2107.00717) to mine these rare examples for labeling."
      ],
      "id": "lOkAr6RnQ2BJ"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Ibaupa-l2T_"
      },
      "source": [
        "# Preparation"
      ],
      "id": "-Ibaupa-l2T_"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ViAJ7rawfp53"
      },
      "source": [
        "## Installation and Imports"
      ],
      "id": "ViAJ7rawfp53"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nmhefeJQfqKb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "462ffa03-b9ce-4a46-da06-0c4698149c04"
      },
      "source": [
        "# Get DISTIL\n",
        "!git clone https://github.com/decile-team/distil.git\n",
        "!pip install -r distil/requirements/requirements.txt\n",
        "\n",
        "# Get MedMNIST\n",
        "!git clone https://github.com/MedMNIST/MedMNIST.git\n",
        "\n",
        "import copy\n",
        "import math\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",
        "sys.path.append(\"MedMNIST/\")\n",
        "from medmnist import PneumoniaMNIST\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 SMI, 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"
      ],
      "id": "nmhefeJQfqKb",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'distil'...\n",
            "remote: Enumerating objects: 3318, done.\u001b[K\n",
            "remote: Counting objects: 100% (1275/1275), done.\u001b[K\n",
            "remote: Compressing objects: 100% (811/811), done.\u001b[K\n",
            "remote: Total 3318 (delta 788), reused 833 (delta 456), pack-reused 2043\u001b[K\n",
            "Receiving objects: 100% (3318/3318), 23.04 MiB | 15.06 MiB/s, done.\n",
            "Resolving deltas: 100% (2061/2061), done.\n",
            "Looking in indexes: https://test.pypi.org/simple/, https://pypi.org/simple/\n",
            "Collecting sphinxcontrib-bibtex>=2.3.0\n",
            "  Downloading sphinxcontrib_bibtex-2.4.1-py3-none-any.whl (38 kB)\n",
            "Collecting multipledispatch==0.6.0\n",
            "  Downloading multipledispatch-0.6.0-py3-none-any.whl (11 kB)\n",
            "Collecting scikit-learn==0.23.0\n",
            "  Downloading scikit_learn-0.23.0-cp37-cp37m-manylinux1_x86_64.whl (7.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 7.3 MB 8.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.14.2 in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 4)) (1.21.5)\n",
            "Requirement already satisfied: scipy>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 5)) (1.4.1)\n",
            "Requirement already satisfied: torch>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 6)) (1.10.0+cu111)\n",
            "Requirement already satisfied: tqdm>=4.24.0 in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 7)) (4.62.3)\n",
            "Requirement already satisfied: numba>=0.43.0 in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 8)) (0.51.2)\n",
            "Collecting submodlib\n",
            "  Downloading https://test-files.pythonhosted.org/packages/55/62/88e02a0e170498f38f7b9ce22b3e0a6a3cf9c82a33d3553da693c5c52872/submodlib-1.1.5.tar.gz (83 kB)\n",
            "\u001b[K     |████████████████████████████████| 83 kB 1.7 MB/s \n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 12)) (1.3.5)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from -r distil/requirements/requirements.txt (line 13)) (0.11.1+cu111)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from multipledispatch==0.6.0->-r distil/requirements/requirements.txt (line 2)) (1.15.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn==0.23.0->-r distil/requirements/requirements.txt (line 3)) (3.1.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn==0.23.0->-r distil/requirements/requirements.txt (line 3)) (1.1.0)\n",
            "Collecting Sphinx>=2.1\n",
            "  Downloading Sphinx-4.4.0-py3-none-any.whl (3.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.1 MB 51.0 MB/s \n",
            "\u001b[?25hCollecting pybtex>=0.20\n",
            "  Downloading pybtex-0.24.0-py2.py3-none-any.whl (561 kB)\n",
            "\u001b[K     |████████████████████████████████| 561 kB 73.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: docutils>=0.8 in /usr/local/lib/python3.7/dist-packages (from sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (0.17.1)\n",
            "Collecting pybtex-docutils>=1.0.0\n",
            "  Downloading pybtex_docutils-1.0.1-py3-none-any.whl (4.8 kB)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.4.0->-r distil/requirements/requirements.txt (line 6)) (3.10.0.2)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from numba>=0.43.0->-r distil/requirements/requirements.txt (line 8)) (57.4.0)\n",
            "Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba>=0.43.0->-r distil/requirements/requirements.txt (line 8)) (0.34.0)\n",
            "Requirement already satisfied: PyYAML>=3.01 in /usr/local/lib/python3.7/dist-packages (from pybtex>=0.20->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (3.13)\n",
            "Collecting latexcodec>=1.0.4\n",
            "  Downloading latexcodec-2.0.1-py2.py3-none-any.whl (18 kB)\n",
            "Collecting sphinxcontrib-htmlhelp>=2.0.0\n",
            "  Downloading sphinxcontrib_htmlhelp-2.0.0-py2.py3-none-any.whl (100 kB)\n",
            "\u001b[K     |████████████████████████████████| 100 kB 13.2 MB/s \n",
            "\u001b[?25hCollecting sphinxcontrib-applehelp\n",
            "  Downloading sphinxcontrib_applehelp-1.0.2-py2.py3-none-any.whl (121 kB)\n",
            "\u001b[K     |████████████████████████████████| 121 kB 79.1 MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (21.3)\n",
            "Requirement already satisfied: sphinxcontrib-serializinghtml>=1.1.5 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (1.1.5)\n",
            "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (4.11.1)\n",
            "Requirement already satisfied: babel>=1.3 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.9.1)\n",
            "Requirement already satisfied: snowballstemmer>=1.1 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.2.0)\n",
            "Collecting sphinxcontrib-devhelp\n",
            "  Downloading sphinxcontrib_devhelp-1.0.2-py2.py3-none-any.whl (84 kB)\n",
            "\u001b[K     |████████████████████████████████| 84 kB 4.8 MB/s \n",
            "\u001b[?25hCollecting sphinxcontrib-jsmath\n",
            "  Downloading sphinxcontrib_jsmath-1.0.1-py2.py3-none-any.whl (5.1 kB)\n",
            "Requirement already satisfied: alabaster<0.8,>=0.7 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (0.7.12)\n",
            "Requirement already satisfied: requests>=2.5.0 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.23.0)\n",
            "Requirement already satisfied: imagesize in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (1.3.0)\n",
            "Collecting sphinxcontrib-qthelp\n",
            "  Downloading sphinxcontrib_qthelp-1.0.3-py2.py3-none-any.whl (90 kB)\n",
            "\u001b[K     |████████████████████████████████| 90 kB 13.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: Jinja2>=2.3 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.11.3)\n",
            "Requirement already satisfied: Pygments>=2.0 in /usr/local/lib/python3.7/dist-packages (from Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.6.1)\n",
            "Requirement already satisfied: pytz>=2015.7 in /usr/local/lib/python3.7/dist-packages (from babel>=1.3->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2018.9)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (3.7.0)\n",
            "Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from Jinja2>=2.3->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.0.1)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.5.0->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.5.0->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.5.0->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (2021.10.8)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.5.0->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (1.24.3)\n",
            "Requirement already satisfied: sklearn in /usr/local/lib/python3.7/dist-packages (from submodlib->-r distil/requirements/requirements.txt (line 11)) (0.0)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas->-r distil/requirements/requirements.txt (line 12)) (2.8.2)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision->-r distil/requirements/requirements.txt (line 13)) (7.1.2)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->Sphinx>=2.1->sphinxcontrib-bibtex>=2.3.0->-r distil/requirements/requirements.txt (line 1)) (3.0.7)\n",
            "Building wheels for collected packages: submodlib\n",
            "  Building wheel for submodlib (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for submodlib: filename=submodlib-1.1.5-cp37-cp37m-linux_x86_64.whl size=491515 sha256=5d44879aa630fba4d0fd08939d3f86fb0714d44b656452cb28ce0a4f77222f52\n",
            "  Stored in directory: /root/.cache/pip/wheels/17/3c/03/48ce7dd03798c0564b61e61020e733443aed88e115518442b4\n",
            "Successfully built submodlib\n",
            "Installing collected packages: latexcodec, sphinxcontrib-qthelp, sphinxcontrib-jsmath, sphinxcontrib-htmlhelp, sphinxcontrib-devhelp, sphinxcontrib-applehelp, scikit-learn, pybtex, Sphinx, pybtex-docutils, submodlib, sphinxcontrib-bibtex, multipledispatch\n",
            "  Attempting uninstall: scikit-learn\n",
            "    Found existing installation: scikit-learn 1.0.2\n",
            "    Uninstalling scikit-learn-1.0.2:\n",
            "      Successfully uninstalled scikit-learn-1.0.2\n",
            "  Attempting uninstall: Sphinx\n",
            "    Found existing installation: Sphinx 1.8.6\n",
            "    Uninstalling Sphinx-1.8.6:\n",
            "      Successfully uninstalled Sphinx-1.8.6\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "imbalanced-learn 0.8.1 requires scikit-learn>=0.24, but you have scikit-learn 0.23.0 which is incompatible.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "Successfully installed Sphinx-4.4.0 latexcodec-2.0.1 multipledispatch-0.6.0 pybtex-0.24.0 pybtex-docutils-1.0.1 scikit-learn-0.23.0 sphinxcontrib-applehelp-1.0.2 sphinxcontrib-bibtex-2.4.1 sphinxcontrib-devhelp-1.0.2 sphinxcontrib-htmlhelp-2.0.0 sphinxcontrib-jsmath-1.0.1 sphinxcontrib-qthelp-1.0.3 submodlib-1.1.5\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "sphinxcontrib"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'MedMNIST'...\n",
            "remote: Enumerating objects: 361, done.\u001b[K\n",
            "remote: Counting objects: 100% (185/185), done.\u001b[K\n",
            "remote: Compressing objects: 100% (125/125), done.\u001b[K\n",
            "remote: Total 361 (delta 110), reused 124 (delta 56), pack-reused 176\u001b[K\n",
            "Receiving objects: 100% (361/361), 5.84 MiB | 12.19 MiB/s, done.\n",
            "Resolving deltas: 100% (203/203), done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v4oIq8U7f6Dw"
      },
      "source": [
        "## Preparing PneumoniaMNIST with Rare Classes\n",
        "\n",
        "The PneumoniaMNIST dataset contains 5,856 32x32 grayscale images in 2 different classes. The 2 classes indicate whether a chest x-ray is indicative of pneumonia or is not indicative of pneumonia. The training set contains 4,708 images, and the test set contains 624 images (the rest are available in a validation set). Here, we do a simple setup of the PneumoniaMNIST dataset that we will use in this example. More importantly, we define a split on PneumoniaMNIST's training set into an initial labeled seed set and an unlabeled set. We also impose an artificial imbalance among the classes to simulate a rare-class scenario. Drawing parallels to actual diagnostics, one may expect to encounter positive cases of pneumonia less frequently than negative cases, showcasing the need for the ability to mine rare classes."
      ],
      "id": "v4oIq8U7f6Dw"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J39lJ2uWVpbM"
      },
      "source": [
        "**Calculate Average/STD**"
      ],
      "id": "J39lJ2uWVpbM"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Khfi12OlSTSz",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 101,
          "referenced_widgets": [
            "4e43acf6a8b74698bba7e8ef07039557",
            "6cba53ea45c5468fafa46f482f97e517",
            "2ca9a395a43e4cc1a7d59a04368c845a",
            "44be0278f4fc47ec8cf438fde4ec12b8",
            "8114758516654ce6972df8a268776717",
            "be7b11692fba4d2593113dfa61c88e4f",
            "e992b3a0110c4292997f0560fd34f87b",
            "4be0dd51017b4b879c5b7d8255a88ae5",
            "732b2a477e594cac82175561e8574237",
            "02cc627dbedf4361804e840016701708",
            "97f8e9a087a94a6eb83975bd9c647f78"
          ]
        },
        "outputId": "63662b5f-b7ee-4dcc-a0b0-adfc34668489"
      },
      "source": [
        "# We do not have the average and standard deviation to use for data normalization. Here, we compute \n",
        "# it on a per-pixel basis, and we use the training set only for this calculation.\n",
        "\n",
        "train_dataset = PneumoniaMNIST(root=\".\", split=\"train\", download=True, transform=transforms.ToTensor())\n",
        "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=False)\n",
        "\n",
        "average = 0.\n",
        "num_pixels = 0\n",
        "for batch_idx, (image, label) in enumerate(train_dataloader):\n",
        "    image = image.to(\"cuda\")\n",
        "    average += image.sum()\n",
        "    num_pixels += len(image.flatten())\n",
        "average = average.item() / num_pixels\n",
        "\n",
        "var = 0.\n",
        "for batch_idx, (image, label) in enumerate(train_dataloader):\n",
        "    image = image.to(\"cuda\").flatten()\n",
        "    image_diff = image - average\n",
        "    image_var = torch.dot(image_diff, image_diff)\n",
        "    var += image_var\n",
        "\n",
        "std = math.sqrt(var / (num_pixels - 1))\n",
        "\n",
        "print(\"AVERAGE:\",average)\n",
        "print(\"STD:\",std)"
      ],
      "id": "Khfi12OlSTSz",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading https://zenodo.org/record/5208230/files/pneumoniamnist.npz?download=1 to ./pneumoniamnist.npz\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "4e43acf6a8b74698bba7e8ef07039557",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "  0%|          | 0/4170669 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "AVERAGE: 0.5719215176512407\n",
            "STD: 0.1683512867181746\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p-8sKY_fBjep"
      },
      "source": [
        "**Create PneumoniaMNIST Dataset**"
      ],
      "id": "p-8sKY_fBjep"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uaaAWcsKf6Uv",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dadb08ae-9636-45e3-f4d5-25cfeab10e2f"
      },
      "source": [
        "# Define the name of the dataset and the path that PyTorch should use when downloading the data\n",
        "data_set_name = 'PneumoniaMNIST'\n",
        "download_path = '.'\n",
        "\n",
        "# Define the number of classes in PneumoniaMNIST\n",
        "nclasses = 2\n",
        "\n",
        "# Define transforms on the dataset splits of PneumoniaMNIST. 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 PneumoniaMNIST.\n",
        "pmnist_training_transform = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((average,), (std,))])\n",
        "pmnist_test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((average,), (std,))])\n",
        "\n",
        "# PneumoniaMNIST provides its labels in list form, so we need to get the original label\n",
        "target_transform = lambda x: x[0]\n",
        "\n",
        "# Get the dataset objects from MedMNIST. Here, PneumoniaMNIST is downloaded, and the transform is applied when points \n",
        "# are retrieved.\n",
        "pmnist_full_train = PneumoniaMNIST(root=download_path, split=\"train\", download=True, transform=pmnist_training_transform, target_transform=target_transform)\n",
        "pmnist_test = PneumoniaMNIST(root=download_path, split=\"test\", download=True, transform=pmnist_test_transform, target_transform=target_transform)\n",
        "\n",
        "# Get the dimension of the images. Here, we simply take the very first image of PneumoniaMNIST\n",
        "# and query its dimension.\n",
        "dim = np.shape(pmnist_full_train[0][0])\n",
        "\n",
        "# We now define a train-unlabeled split for the sake of the experiment. Here, we specify the ratio of common classes to rare classes\n",
        "# for both the initial training seed set and the unlabeled set. We use a ratio of 100:1.\n",
        "num_examples_per_common_class_seed = 200\n",
        "num_examples_per_rare_class_seed = 2\n",
        "\n",
        "num_examples_per_common_class_unlabeled = 1000\n",
        "num_examples_per_rare_class_unlabeled = 10\n",
        "\n",
        "# We create the imbalance on class 1.\n",
        "rare_classes = [1]\n",
        "\n",
        "# Create the imbalance by choosing which indices of the full training set to assign to the initial labeled seed set and the unlabeled set\n",
        "train_idx = []\n",
        "unlabeled_idx = []\n",
        "\n",
        "for class_idx in range(nclasses):\n",
        "\n",
        "    # Retrieve all the indices of the elements in PneumoniaMNIST whose label matches class_idx\n",
        "    full_idx_class = torch.where(torch.Tensor(pmnist_full_train.labels) == class_idx)[0].cpu().numpy()\n",
        "\n",
        "    # Determine how many points to add for this class depending on the rarity of the class\n",
        "    if class_idx in rare_classes:\n",
        "        class_num_training_examples_to_add = num_examples_per_rare_class_seed\n",
        "        class_num_unlabeled_examples_to_add = num_examples_per_rare_class_unlabeled\n",
        "    else:\n",
        "        class_num_training_examples_to_add = num_examples_per_common_class_seed\n",
        "        class_num_unlabeled_examples_to_add = num_examples_per_common_class_unlabeled\n",
        "\n",
        "    # Choose randomly a subset of these indices. These will be added to the initial training seed set.\n",
        "    class_train_idx = np.random.choice(full_idx_class, size=class_num_training_examples_to_add, replace=False)\n",
        "\n",
        "    # Choose randomly a subset of the remaining indices. These will be added to the unlabeled set.\n",
        "    remaining_class_idx = np.array(list(set(full_idx_class) - set(class_train_idx)))\n",
        "    class_unlabeled_idx = np.random.choice(remaining_class_idx, size=class_num_unlabeled_examples_to_add, replace=False)\n",
        "\n",
        "    # Add the chosen indices to the growing subset lists\n",
        "    train_idx.extend(class_train_idx)\n",
        "    unlabeled_idx.extend(class_unlabeled_idx)\n",
        "\n",
        "# Create the train and unlabeled subsets based on the index lists above. 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",
        "pmnist_train = Subset(pmnist_full_train, train_idx)\n",
        "pmnist_unlabeled = Subset(pmnist_full_train, unlabeled_idx)"
      ],
      "id": "uaaAWcsKf6Uv",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using downloaded and verified file: ./pneumoniamnist.npz\n",
            "Using downloaded and verified file: ./pneumoniamnist.npz\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mjVjSm-slI7U"
      },
      "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. In this case, we need a two-class output and a one-channel input."
      ],
      "id": "mjVjSm-slI7U"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LjHfyJvslJID"
      },
      "source": [
        "net = ResNet18(num_classes=nclasses, channels=1)"
      ],
      "id": "LjHfyJvslJID",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2fsSP89pl7Rk"
      },
      "source": [
        "# Initial Round"
      ],
      "id": "2fsSP89pl7Rk"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eGAnKw8AoE_q"
      },
      "source": [
        "## Training\n",
        "\n",
        "Now that we have prepared the training data and the model, we can begin training an initial model. We use DISTIL's provided [training loop](https://github.com/decile-team/distil/blob/main/distil/utils/train_helper.py) to do training."
      ],
      "id": "eGAnKw8AoE_q"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f6_PV51rl7f0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e0c29fb8-def1-4edc-ed0d-a582cf10509e"
      },
      "source": [
        "# Define the training arguments to use.\n",
        "args = {'n_epoch':      300,    # Stop training after 300 epochs.\n",
        "        'lr':           0.01,   # Use a learning rate of 0.01\n",
        "        'batch_size':   20,     # Update the parameters using training batches of size 20\n",
        "        'max_accuracy': 0.99,   # Stop training once the training accuracy has exceeded 0.99\n",
        "        'optimizer':    'sgd',  # Use the stochastic gradient descent optimizer\n",
        "        'device':       \"cuda\" if torch.cuda.is_available() else \"cpu\"  # Use a GPU if one is available\n",
        "        }\n",
        "\n",
        "# Create the training loop using our training dataset, provided model, and training arguments.\n",
        "# Train an initial model.\n",
        "dt = data_train(pmnist_train, net, args)\n",
        "trained_model = dt.train()"
      ],
      "id": "f6_PV51rl7f0",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training..\n",
            "Epoch: 3 Training accuracy: 0.99\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yY13YA4GoIvY"
      },
      "source": [
        "## Evaluation\n",
        "\n",
        "How does our initial model do on the test set? Furthermore, how does our initial model do on the *rare classes* of the test set? Luckily, the training loop provided by DISTIL also provides a way to measure the accuracy of the model on a given dataset. We measure both accuracies here."
      ],
      "id": "yY13YA4GoIvY"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uGgL9oagoqKT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f9e8395f-e16f-4916-841f-a48a370af90c"
      },
      "source": [
        "# Get the full test accuracy\n",
        "full_test_accuracy = dt.get_acc_on_set(pmnist_test)\n",
        "print(F\"Full Test Accuracy: {full_test_accuracy}\")\n",
        "\n",
        "# Get the per-class test accuracies\n",
        "rare_indices = []\n",
        "for class_idx in range(nclasses):\n",
        "\n",
        "    # Get the indices of the test set corresponding to this class\n",
        "    test_rare_class_subset_idx = torch.where(torch.Tensor(pmnist_test.labels) == class_idx)[0].cpu().numpy()\n",
        "    \n",
        "    if class_idx in rare_classes:\n",
        "        rare_indices.extend(test_rare_class_subset_idx)\n",
        "\n",
        "    # Get the accuracy on this class subset\n",
        "    pmnist_test_class_subset = Subset(pmnist_test, test_rare_class_subset_idx)\n",
        "    rare_class_test_accuracy = dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    print(F\"Class {class_idx} Test Accuracy: {rare_class_test_accuracy}\")\n",
        "\n",
        "# Get accuracy on all rare points\n",
        "pmnist_test_rare_subset = Subset(pmnist_test, rare_indices)\n",
        "rare_test_accuracy = dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "print(F\"Rare Test Accuracy: {rare_test_accuracy}\")"
      ],
      "id": "uGgL9oagoqKT",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Full Test Accuracy: 0.375\n",
            "Class 0 Test Accuracy: 1.0\n",
            "Class 1 Test Accuracy: 0.0\n",
            "Rare Test Accuracy: 0.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p-47FnQRsfkM"
      },
      "source": [
        "As we can see, the test performance on the rare classes is awful. Can we rectify this issue by adding rare-class examples?"
      ],
      "id": "p-47FnQRsfkM"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JnDJvyMWtCCa"
      },
      "source": [
        "# Mining Rare Classes"
      ],
      "id": "JnDJvyMWtCCa"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lH1qtKB4tK2v"
      },
      "source": [
        "## Preparing a Query Set\n",
        "\n",
        "In this example, we know that there are rare classes in the unlabeled dataset that can help us improve the accuracy of the model. How do we select these points if we do not know *where* they are in the unlabeled dataset? Here, we use [SIMILAR](https://arxiv.org/abs/2107.00717) to select rare-class examples from the unlabeled set. SIMILAR requires access to a query set of points that it uses to choose similar unlabeled points. Hence, we must first prepare this query set.\n",
        "\n",
        "Where do we get the query set? Luckily, we already have a couple points in our training dataset to choose. We could have otherwise used some held-out points."
      ],
      "id": "lH1qtKB4tK2v"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MQupSdbrur3_"
      },
      "source": [
        "# Go over the training dataset, getting the indices of all the rare-class examples\n",
        "rare_training_example_indices = []\n",
        "for index, (_, label) in enumerate(pmnist_train):\n",
        "    if label in rare_classes:\n",
        "        rare_training_example_indices.append(index)\n",
        "\n",
        "# Create a query set that contains only the rare-class examples of the training dataset\n",
        "rare_class_query_set = Subset(pmnist_train, rare_training_example_indices)"
      ],
      "id": "MQupSdbrur3_",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cgvh20tYviAr"
      },
      "source": [
        "## Using DISTIL's SIMILAR Implementation: SMI\n",
        "\n",
        "Now that we have the query set, we are ready to use DISTIL's implementation of SIMILAR. In particular, we use the submodular mutual information [strategy](https://github.com/decile-team/distil/blob/main/distil/active_learning_strategies/smi.py) that is detailed in [SIMILAR](https://arxiv.org/abs/2107.00717). This will allow us to select a set of points to label within a specified budget $k$.\n",
        "\n",
        "Specifically, the strategy attempts to maximize the [submodular mutual information](https://arxiv.org/abs/2006.15412) between a subset $\\mathcal{A}$ of size no greater than $k$ of the unlabeled dataset $\\mathcal{U}$ and the query set $\\mathcal{Q}$:\n",
        "\n",
        "\\begin{align}\n",
        "\\text{argmax}_{\\mathcal{A} \\subseteq \\mathcal{U}, |\\mathcal{A}|\\leq k} I_F(\\mathcal{A};\\mathcal{Q})\n",
        "\\end{align}\n",
        "\n",
        "where $F$ is a submodular set function."
      ],
      "id": "cgvh20tYviAr"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HPKl0J3HyZZO"
      },
      "source": [
        "# Define arguments for SMI\n",
        "selection_strategy_args = {'device':        args['device'],       # Use the device used in training\n",
        "                           'batch_size':    args['batch_size'],   # Use the batch size used in training\n",
        "                           'smi_function':  'fl2mi',              # Use a facility location function, which captures representation information\n",
        "                           'metric':        'cosine',             # Use cosine similarity when determining the likeness of two data points\n",
        "                           'optimizer':     'LazyGreedy'          # When doing submodular maximization, use the lazy greedy optimizer\n",
        "                          }\n",
        "\n",
        "# Create the SMI selection strategy. Note: We remove the labels from the unlabeled portion of PneumoniaMNIST that we created earlier.\n",
        "# In a practical application, one would not have these labels a priori.\n",
        "selection_strategy = SMI(pmnist_train, LabeledToUnlabeledDataset(pmnist_unlabeled), rare_class_query_set, trained_model, nclasses, selection_strategy_args)\n",
        "\n",
        "# Disable the augmentations used in the training dataset. Since all augmentations come from the pmnist_full_train object, we set its transform to the \n",
        "# transform used by the test set.\n",
        "pmnist_full_train.transform = pmnist_test_transform\n",
        "\n",
        "# Do the selection, which will return the indices of the selected points with respect to the unlabeled dataset.\n",
        "budget = 20\n",
        "selected_idx = selection_strategy.select(budget)\n",
        "\n",
        "# Re-enable augmentations\n",
        "pmnist_full_train.transform = pmnist_training_transform"
      ],
      "id": "HPKl0J3HyZZO",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zoQAIxhA1dLE"
      },
      "source": [
        "## Labeling the Points\n",
        "\n",
        "Now that we know which points should be labeled, we can present them to human labelers for annotation. Here, we can do so automatically since we already know their labels for the sake of the example."
      ],
      "id": "zoQAIxhA1dLE"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HAlQ5Fr51yS8"
      },
      "source": [
        "# Form a labeled subset of the unlabeled dataset. Again, we already have the labels, \n",
        "# so we simply take a subset. Note, however, that the selection was done without the \n",
        "# use of the labels and that we would normally not have these labels. Hence, the \n",
        "# following statement would usually require human effort to complete.\n",
        "smi_human_labeled_dataset = Subset(pmnist_unlabeled, selected_idx)"
      ],
      "id": "HAlQ5Fr51yS8",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zom1rZrJ2Wi_"
      },
      "source": [
        "## Characterizing the Selection\n",
        "\n",
        "Now that we have selected and labeled these new points, we can add them to the training dataset and retrain our model. Before that, how many points did we select that actually were rare-class points?"
      ],
      "id": "Zom1rZrJ2Wi_"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GN22GOF-2W1M",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4b20462a-c98a-4c3a-ab92-01e9cc164dd1"
      },
      "source": [
        "# Go over the newly labeled dataset, tallying the number of points seen in each class.\n",
        "smi_class_counts = [0 for x in range(nclasses)]\n",
        "for _, label in smi_human_labeled_dataset:\n",
        "    smi_class_counts[label] += 1\n",
        "\n",
        "# Print each class count\n",
        "for class_idx, class_count in enumerate(smi_class_counts):\n",
        "    print(F\"Class {class_idx} count: {class_count}\")\n",
        "\n",
        "# Print total rare count\n",
        "total_rare_count = sum([smi_class_counts[i] for i in rare_classes])\n",
        "print(F\"Total Rare Count: {total_rare_count}\")"
      ],
      "id": "GN22GOF-2W1M",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Class 0 count: 10\n",
            "Class 1 count: 10\n",
            "Total Rare Count: 10\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xyZx_yx07sEw"
      },
      "source": [
        "We were able to get a good number of rare-class points. For comparison sake, how many rare points would we get using some of DISTIL's other strategies?"
      ],
      "id": "xyZx_yx07sEw"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FvDNvY3D9hSn"
      },
      "source": [
        "**BADGE**"
      ],
      "id": "FvDNvY3D9hSn"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HN_iPljF9e_L",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "db6f4b49-6362-4eb4-d5e2-cbd40a1414f9"
      },
      "source": [
        "# Repeat the previous steps\n",
        "selection_strategy = BADGE(pmnist_train, LabeledToUnlabeledDataset(pmnist_unlabeled), trained_model, nclasses, selection_strategy_args)\n",
        "\n",
        "pmnist_full_train.transform = pmnist_test_transform\n",
        "budget = 20\n",
        "selected_idx = selection_strategy.select(budget)\n",
        "pmnist_full_train.transform = pmnist_training_transform\n",
        "\n",
        "badge_human_labeled_dataset = Subset(pmnist_unlabeled, selected_idx)\n",
        "\n",
        "badge_class_counts = [0 for x in range(nclasses)]\n",
        "for _, label in badge_human_labeled_dataset:\n",
        "    badge_class_counts[label] += 1\n",
        "\n",
        "for class_idx, class_count in enumerate(badge_class_counts):\n",
        "    print(F\"Class {class_idx} count: {class_count}\")\n",
        "\n",
        "total_rare_count = sum([badge_class_counts[i] for i in rare_classes])\n",
        "print(F\"Total Rare Count: {total_rare_count}\")"
      ],
      "id": "HN_iPljF9e_L",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Class 0 count: 19\n",
            "Class 1 count: 1\n",
            "Total Rare Count: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E1PDZ2qk-hi7"
      },
      "source": [
        "**Random**"
      ],
      "id": "E1PDZ2qk-hi7"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o4DQ3_ld-j8U",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "67f53fab-ac48-46cb-9038-40fac2c00a41"
      },
      "source": [
        "# Repeat the previous steps\n",
        "selection_strategy = RandomSampling(pmnist_train, LabeledToUnlabeledDataset(pmnist_unlabeled), trained_model, nclasses, selection_strategy_args)\n",
        "\n",
        "pmnist_full_train.transform = pmnist_test_transform\n",
        "budget = 20\n",
        "selected_idx = selection_strategy.select(budget)\n",
        "pmnist_full_train.transform = pmnist_training_transform\n",
        "\n",
        "random_human_labeled_dataset = Subset(pmnist_unlabeled, selected_idx)\n",
        "\n",
        "random_class_counts = [0 for x in range(nclasses)]\n",
        "for _, label in random_human_labeled_dataset:\n",
        "    random_class_counts[label] += 1\n",
        "\n",
        "for class_idx, class_count in enumerate(random_class_counts):\n",
        "    print(F\"Class {class_idx} count: {class_count}\")\n",
        "\n",
        "total_rare_count = sum([random_class_counts[i] for i in rare_classes])\n",
        "print(F\"Total Rare Count: {total_rare_count}\")"
      ],
      "id": "o4DQ3_ld-j8U",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Class 0 count: 19\n",
            "Class 1 count: 1\n",
            "Total Rare Count: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wK5KB_UH-zR2"
      },
      "source": [
        "**Entropy**"
      ],
      "id": "wK5KB_UH-zR2"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ncA_-JMA-zoA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9239a45d-6c5e-45ab-b934-06ed376f2af2"
      },
      "source": [
        "# Repeat the previous steps\n",
        "selection_strategy = EntropySampling(pmnist_train, LabeledToUnlabeledDataset(pmnist_unlabeled), trained_model, nclasses, selection_strategy_args)\n",
        "\n",
        "pmnist_full_train.transform = pmnist_test_transform\n",
        "budget = 20\n",
        "selected_idx = selection_strategy.select(budget)\n",
        "pmnist_full_train.transform = pmnist_training_transform\n",
        "\n",
        "entropy_human_labeled_dataset = Subset(pmnist_unlabeled, selected_idx)\n",
        "\n",
        "entropy_class_counts = [0 for x in range(nclasses)]\n",
        "for _, label in entropy_human_labeled_dataset:\n",
        "    entropy_class_counts[label] += 1\n",
        "\n",
        "for class_idx, class_count in enumerate(entropy_class_counts):\n",
        "    print(F\"Class {class_idx} count: {class_count}\")\n",
        "\n",
        "total_rare_count = sum([entropy_class_counts[i] for i in rare_classes])\n",
        "print(F\"Total Rare Count: {total_rare_count}\")"
      ],
      "id": "ncA_-JMA-zoA",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Class 0 count: 19\n",
            "Class 1 count: 1\n",
            "Total Rare Count: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RvMg93--Byvl"
      },
      "source": [
        "Hence, we can see that SMI does comparatively much better at selecting rare instances."
      ],
      "id": "RvMg93--Byvl"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vjvhIDHsAkU1"
      },
      "source": [
        "# Improving Performance"
      ],
      "id": "vjvhIDHsAkU1"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eF-HsA27AsUG"
      },
      "source": [
        "## Re-Training\n",
        "\n",
        "Let us re-train our model using the newly selected points."
      ],
      "id": "eF-HsA27AsUG"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1kkbzgxVAkl7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "302f41ff-f9bd-4cc7-8049-9d4c2de216ee"
      },
      "source": [
        "# Create a new training dataset by concatenating what we have with the newly labeled points\n",
        "new_training_dataset = ConcatDataset([pmnist_train, smi_human_labeled_dataset])\n",
        "new_dt = data_train(new_training_dataset, copy.deepcopy(net), args)\n",
        "new_trained_model = new_dt.train()"
      ],
      "id": "1kkbzgxVAkl7",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training..\n",
            "Epoch: 19 Training accuracy: 1.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6PtUxT2pBLx_"
      },
      "source": [
        "## Evaluation\n",
        "\n",
        "Now, let us see the accuracy improvement."
      ],
      "id": "6PtUxT2pBLx_"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Xj8wmgPBP9U",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "24cdab24-bdcd-463e-8100-ee3ec38db14c"
      },
      "source": [
        "# Get the full test accuracy\n",
        "full_test_accuracy_before = dt.get_acc_on_set(pmnist_test)\n",
        "full_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test)\n",
        "\n",
        "print(F\"Full Test Accuracy Improvement: {full_test_accuracy_before} to {full_test_accuracy_after}\")\n",
        "\n",
        "# Get the per-class test accuracies\n",
        "rare_indices = []\n",
        "for class_idx in range(nclasses):\n",
        "\n",
        "    # Get the indices of the test set corresponding to this class\n",
        "    test_rare_class_subset_idx = torch.where(torch.Tensor(pmnist_test.labels) == class_idx)[0].cpu().numpy()\n",
        "    \n",
        "    if class_idx in rare_classes:\n",
        "        rare_indices.extend(test_rare_class_subset_idx)\n",
        "\n",
        "    # Get the accuracy on this class subset\n",
        "    pmnist_test_class_subset = Subset(pmnist_test, test_rare_class_subset_idx)\n",
        "\n",
        "    rare_class_test_accuracy_before = dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    rare_class_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    print(F\"Class {class_idx} Test Accuracy: {rare_class_test_accuracy_before} to {rare_class_test_accuracy_after}\")\n",
        "\n",
        "# Get accuracy on all rare points\n",
        "pmnist_test_rare_subset = Subset(pmnist_test, rare_indices)\n",
        "rare_test_accuracy_before = dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "rare_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "print(F\"Rare Test Accuracy: {rare_test_accuracy_before} to {rare_test_accuracy_after}\")"
      ],
      "id": "2Xj8wmgPBP9U",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Full Test Accuracy Improvement: 0.375 to 0.7900641025641025\n",
            "Class 0 Test Accuracy: 1.0 to 0.9487179487179487\n",
            "Class 1 Test Accuracy: 0.0 to 0.6948717948717948\n",
            "Rare Test Accuracy: 0.0 to 0.6948717948717948\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BK89JhX6Fq3s"
      },
      "source": [
        "## Comparison\n",
        "\n",
        "What would the accuracy improvement look like if we had used the other methods?"
      ],
      "id": "BK89JhX6Fq3s"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jxklRX4uF1UO"
      },
      "source": [
        "**BADGE**"
      ],
      "id": "jxklRX4uF1UO"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X0_1FoFYF8pU",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b9c0ccb6-5e14-4dc2-8ade-7eae97c731a1"
      },
      "source": [
        "# Repeat the process\n",
        "new_training_dataset = ConcatDataset([pmnist_train, badge_human_labeled_dataset])\n",
        "new_dt = data_train(new_training_dataset, copy.deepcopy(net), args)\n",
        "new_trained_model = new_dt.train()\n",
        "\n",
        "full_test_accuracy_before = dt.get_acc_on_set(pmnist_test)\n",
        "full_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test)\n",
        "\n",
        "print(F\"Full Test Accuracy Improvement: {full_test_accuracy_before} to {full_test_accuracy_after}\")\n",
        "\n",
        "rare_indices = []\n",
        "for class_idx in range(nclasses):\n",
        "\n",
        "    test_rare_class_subset_idx = torch.where(torch.Tensor(pmnist_test.labels) == class_idx)[0].cpu().numpy()\n",
        "    \n",
        "    if class_idx in rare_classes:\n",
        "        rare_indices.extend(test_rare_class_subset_idx)\n",
        "\n",
        "    pmnist_test_class_subset = Subset(pmnist_test, test_rare_class_subset_idx)\n",
        "\n",
        "    rare_class_test_accuracy_before = dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    rare_class_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    print(F\"Class {class_idx} Test Accuracy: {rare_class_test_accuracy_before} to {rare_class_test_accuracy_after}\")\n",
        "\n",
        "pmnist_test_rare_subset = Subset(pmnist_test, rare_indices)\n",
        "rare_test_accuracy_before = dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "rare_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "print(F\"Rare Test Accuracy: {rare_test_accuracy_before} to {rare_test_accuracy_after}\")"
      ],
      "id": "X0_1FoFYF8pU",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training..\n",
            "Epoch: 8 Training accuracy: 0.991\n",
            "Full Test Accuracy Improvement: 0.375 to 0.375\n",
            "Class 0 Test Accuracy: 1.0 to 1.0\n",
            "Class 1 Test Accuracy: 0.0 to 0.0\n",
            "Rare Test Accuracy: 0.0 to 0.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gNArqgcGF41l"
      },
      "source": [
        "**Random**"
      ],
      "id": "gNArqgcGF41l"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xUgD2s2lF86O",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5ac9a9b4-d485-4251-ad1e-d75ef2bb7cd2"
      },
      "source": [
        "# Repeat the process\n",
        "new_training_dataset = ConcatDataset([pmnist_train, random_human_labeled_dataset])\n",
        "new_dt = data_train(new_training_dataset, copy.deepcopy(net), args)\n",
        "new_trained_model = new_dt.train()\n",
        "\n",
        "full_test_accuracy_before = dt.get_acc_on_set(pmnist_test)\n",
        "full_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test)\n",
        "\n",
        "print(F\"Full Test Accuracy Improvement: {full_test_accuracy_before} to {full_test_accuracy_after}\")\n",
        "\n",
        "rare_indices = []\n",
        "for class_idx in range(nclasses):\n",
        "\n",
        "    test_rare_class_subset_idx = torch.where(torch.Tensor(pmnist_test.labels) == class_idx)[0].cpu().numpy()\n",
        "    \n",
        "    if class_idx in rare_classes:\n",
        "        rare_indices.extend(test_rare_class_subset_idx)\n",
        "\n",
        "    pmnist_test_class_subset = Subset(pmnist_test, test_rare_class_subset_idx)\n",
        "\n",
        "    rare_class_test_accuracy_before = dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    rare_class_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    print(F\"Class {class_idx} Test Accuracy: {rare_class_test_accuracy_before} to {rare_class_test_accuracy_after}\")\n",
        "\n",
        "pmnist_test_rare_subset = Subset(pmnist_test, rare_indices)\n",
        "rare_test_accuracy_before = dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "rare_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "print(F\"Rare Test Accuracy: {rare_test_accuracy_before} to {rare_test_accuracy_after}\")"
      ],
      "id": "xUgD2s2lF86O",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training..\n",
            "Epoch: 12 Training accuracy: 0.986\n",
            "Full Test Accuracy Improvement: 0.375 to 0.375\n",
            "Class 0 Test Accuracy: 1.0 to 1.0\n",
            "Class 1 Test Accuracy: 0.0 to 0.0\n",
            "Rare Test Accuracy: 0.0 to 0.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SkcyKstqF5Ml"
      },
      "source": [
        "**Entropy**"
      ],
      "id": "SkcyKstqF5Ml"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nlB34d8fF9bP",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "016029f9-bab4-4d91-9140-683351a0e18f"
      },
      "source": [
        "# Repeat the process\n",
        "new_training_dataset = ConcatDataset([pmnist_train, entropy_human_labeled_dataset])\n",
        "new_dt = data_train(new_training_dataset, copy.deepcopy(net), args)\n",
        "new_trained_model = new_dt.train()\n",
        "\n",
        "full_test_accuracy_before = dt.get_acc_on_set(pmnist_test)\n",
        "full_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test)\n",
        "\n",
        "print(F\"Full Test Accuracy Improvement: {full_test_accuracy_before} to {full_test_accuracy_after}\")\n",
        "\n",
        "rare_indices = []\n",
        "for class_idx in range(nclasses):\n",
        "\n",
        "    test_rare_class_subset_idx = torch.where(torch.Tensor(pmnist_test.labels) == class_idx)[0].cpu().numpy()\n",
        "    \n",
        "    if class_idx in rare_classes:\n",
        "        rare_indices.extend(test_rare_class_subset_idx)\n",
        "\n",
        "    pmnist_test_class_subset = Subset(pmnist_test, test_rare_class_subset_idx)\n",
        "\n",
        "    rare_class_test_accuracy_before = dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    rare_class_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_class_subset)\n",
        "    print(F\"Class {class_idx} Test Accuracy: {rare_class_test_accuracy_before} to {rare_class_test_accuracy_after}\")\n",
        "\n",
        "pmnist_test_rare_subset = Subset(pmnist_test, rare_indices)\n",
        "rare_test_accuracy_before = dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "rare_test_accuracy_after = new_dt.get_acc_on_set(pmnist_test_rare_subset)\n",
        "print(F\"Rare Test Accuracy: {rare_test_accuracy_before} to {rare_test_accuracy_after}\")"
      ],
      "id": "nlB34d8fF9bP",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training..\n",
            "Epoch: 9 Training accuracy: 0.991\n",
            "Full Test Accuracy Improvement: 0.375 to 0.375\n",
            "Class 0 Test Accuracy: 1.0 to 1.0\n",
            "Class 1 Test Accuracy: 0.0 to 0.0\n",
            "Rare Test Accuracy: 0.0 to 0.0\n"
          ]
        }
      ]
    }
  ]
}