project_code.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# numpy to manipulate and make computations on arrays\n",
    "import numpy as np\n",
    "\n",
    "# libraries to process datasets\n",
    "#from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "# importing evaluation metrics and splitting function from sklearn \n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score, confusion_matrix, classification_report\n",
    "\n",
    "# importing PyTorch libraries\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision.transforms import v2\n",
    "\n",
    "# wrapper allowing to use PyTorch with sklearn\n",
    "from skorch import NeuralNetClassifier\n",
    "\n",
    "# Active learning library\n",
    "from collections.abc import Mapping\n",
    "from typing import Callable\n",
    "from sklearn.base import BaseEstimator\n",
    "from modAL.utils.data import modALinput\n",
    "from modAL.utils.selection import multi_argmax, shuffled_argmax\n",
    "from modAL.dropout import default_logits_adaptor, set_dropout_mode, _bald_divergence\n",
    "from modAL.models import DeepActiveLearner\n",
    "\n",
    "# libraries for plotting\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "# import wandb for experiment tracking and login\n",
    "# instructions to wandb login here: <https://docs.wandb.ai/quickstart#:~:text=Sign%20up%20for%20a%20free,Python%203%20environment%20using%20pip%20>\n",
    "import wandb\n",
    "wandb.login()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download whole dataset in numpy format \n",
    "# change path according to your location\n",
    "# load datasets\n",
    "train_data = np.load(\"kmnist-master/dataset/kmnist-train-imgs.npz\")['arr_0']\n",
    "train_labels = np.load(\"kmnist-master/dataset/kmnist-train-labels.npz\")['arr_0']\n",
    "test_data = np.load(\"kmnist-master/dataset/kmnist-test-imgs.npz\")['arr_0']\n",
    "test_labels = np.load(\"kmnist-master/dataset/kmnist-test-labels.npz\")['arr_0']\n",
    "\n",
    "num_class = len(np.unique(train_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split training set to have also a validation set\n",
    "\n",
    "train_data, val_data, train_labels, val_labels = train_test_split(train_data, train_labels, test_size = 0.2, random_state = 42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# It's already available online for this dataset and it's a default value, but to double check\n",
    "# Get mean and std of pixel values in dataset\n",
    "# Then divide those by 255 to normalize values to range [0,1]\n",
    "# 255 because pixel values range from 0 to 255. 0 representing black and 255 representing white\n",
    "\n",
    "print(train_data.shape)\n",
    "mean = np.mean(train_data)\n",
    "print(mean/255)\n",
    "std = np.std(train_data)\n",
    "print(std/255)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize example\n",
    "ind = np.random.randint(0, len(train_data)-1)\n",
    "plt.imshow(train_data[ind].reshape((28, 28)), cmap='gray')\n",
    "plt.title(\"Label: {}\".format(train_labels[ind]))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset class for handling KMNIST, consisting of images and corresponding labels, with optional transformations\n",
    "\n",
    "class KMNISTDataset(Dataset):\n",
    "    def __init__(self, images, labels, transform=None):\n",
    "        self.images = images\n",
    "        self.labels = labels\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "\n",
    "    def __getitem__(self, n):\n",
    "        image = self.images[n].reshape((28, 28)).astype(np.uint8)\n",
    "        label = self.labels[n]\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return (image, label)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = range(10)\n",
    "\n",
    "# Define transformations that will convert and normalize data\n",
    "\n",
    "# Training data has also been augmented \n",
    "train_transform = v2.Compose([\n",
    "    v2.ToPILImage(),\n",
    "    v2.PILToTensor(),\n",
    "    v2.ToDtype(torch.float32, scale=True),\n",
    "    v2.RandomAffine(degrees=20, translate=(0.1,0.1), scale=(0.9, 1.1)),\n",
    "    v2.ColorJitter(brightness=0.2, contrast=0.2),\n",
    "    v2.Normalize(mean=[0.1917], std=[0.3483])\n",
    "])\n",
    "\n",
    "val_normalization = v2.Compose([\n",
    "    v2.ToPILImage(),\n",
    "    v2.PILToTensor(),\n",
    "    v2.ToDtype(torch.float32, scale=True),\n",
    "    v2.Normalize(mean=[0.1917], std=[0.3483])\n",
    "])\n",
    "\n",
    "test_normalization = val_normalization\n",
    "\n",
    "# Apply transformations and load datasets\n",
    "\n",
    "train_dataset = KMNISTDataset(train_data, train_labels, transform = train_transform)\n",
    "\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=48000, shuffle = True)\n",
    "\n",
    "val_dataset = KMNISTDataset(val_data, val_labels, transform = val_normalization)\n",
    "\n",
    "val_loader = DataLoader(dataset=val_dataset, batch_size=12000,shuffle = True)\n",
    "\n",
    "test_dataset = KMNISTDataset(test_data, test_labels, transform = test_normalization)\n",
    "\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=12000,shuffle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A series of sanity checks will follow. This is to make sure data has been preprocessed well. \n",
    "\n",
    "Also, next(iter()) is called to iterate over batches of data\n",
    "\n",
    "Normalization should give a mean close to 0 and a standard deviation close to 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check for training data\n",
    "X_train, y_train = next(iter(train_loader))\n",
    "\n",
    "X_train[0].data.shape\n",
    "print(X_train.data.min())\n",
    "print(X_train.data.max())\n",
    "# after normalization, mean should be around 0 and std should be close to 1\n",
    "print(\"Mean:\", X_train.data.mean())\n",
    "print(\"Std:\", X_train.data.std())\n",
    "print(\"Label:\", classes[y_train[0]])\n",
    "plt.imshow(X_train[0].data.reshape((28,28)), cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check for validation data\n",
    "X_val, y_val = next(iter(val_loader))\n",
    "\n",
    "X_val[0].data.shape\n",
    "print(X_val.data.min())\n",
    "print(X_val.data.max())\n",
    "print(\"Mean:\", X_val.data.mean())\n",
    "print(\"Std:\", X_val.data.std())\n",
    "print(\"Label:\", classes[y_val[0]])\n",
    "plt.imshow(X_val[0].data.reshape((28,28)), cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check for test data\n",
    "X_test, y_test = next(iter(test_loader))\n",
    "\n",
    "X_test[0].data.shape\n",
    "print(X_test.data.min())\n",
    "print(X_test.data.max())\n",
    "print(\"Mean:\", X_test.data.mean())\n",
    "print(\"Std:\", X_test.data.std())\n",
    "print(\"Label:\", classes[y_test[0]])\n",
    "plt.imshow(X_test[0].data.reshape((28,28)), cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another sanity check - showing 5 samples from three sub-datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to show data\n",
    "def show_images(images, labels, title=\"Images\"):\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for i in range(min(5, len(images))):  # Display up to 5 images\n",
    "        plt.subplot(1, 5, i + 1)\n",
    "        plt.imshow(images[i].squeeze(), cmap='gray')\n",
    "        plt.title(f\"Label: {labels[i]}\")\n",
    "        plt.axis('off')\n",
    "    plt.suptitle(title)\n",
    "    plt.show()\n",
    "\n",
    "# Convert tensor to numpy array\n",
    "X_train_numpy = X_train.numpy()\n",
    "X_val_numpy = X_val.numpy()\n",
    "X_test_numpy = X_test.numpy()\n",
    "\n",
    "# Show samples from the training set\n",
    "show_images(X_train_numpy, y_train, title=\"Training Set Samples\")\n",
    "\n",
    "# Show samples from the validation set\n",
    "show_images(X_val_numpy, y_val, title=\"Validation Set Samples\")\n",
    "\n",
    "# Show samples from the test set\n",
    "show_images(X_test_numpy, y_test, title=\"Test Set Samples\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assemble initial data and generate pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# assemble initial data\n",
    "n_initial = 1000\n",
    "initial_idx = np.random.choice(range(len(X_train)), size=n_initial, replace=False)\n",
    "X_initial = X_train[initial_idx]\n",
    "y_initial = y_train[initial_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate the pool\n",
    "# remove the initial data from the training dataset\n",
    "X_pool = np.delete(X_train, initial_idx, axis=0)\n",
    "y_pool = np.delete(y_train, initial_idx, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNNmodel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNNmodel, self).__init__()\n",
    "        \n",
    "        self.conv1 = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.Dropout2d(0.4),\n",
    "            nn.Conv2d(32, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.Dropout2d(0.4),\n",
    "            nn.Conv2d(32, 32, 3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Dropout2d(0.4)\n",
    "        )\n",
    "        \n",
    "        self.conv2 = nn.Sequential(\n",
    "            nn.Conv2d(32, 64, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.Dropout2d(0.4),\n",
    "            nn.Conv2d(64, 64, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.Dropout2d(0.4),\n",
    "            nn.Conv2d(64, 64, 3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Dropout2d(0.4)\n",
    "        )\n",
    "        \n",
    "        self.conv3 = nn.Sequential(\n",
    "            nn.Conv2d(64, 128, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Dropout2d(0.4)\n",
    "        )\n",
    "        \n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(128, 10)\n",
    "        )\n",
    "                \n",
    "        \n",
    "        self.input_dropout = nn.Dropout(0.7)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.input_dropout(x)\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.conv3(x)\n",
    "        \n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x)\n",
    "        x = F.log_softmax(x, dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an instance of model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CNNmodel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get overview of model\n",
    "from torchsummary import summary\n",
    "summary(model, (1, 28, 28))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set device to gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set gpu. Note: mps is Apple M1's \"cuda\". Please change mps to cuda if you run this code on a non-Apple chip\n",
    "\n",
    "# Make sure mps is available\n",
    "torch.backends.mps.is_available()\n",
    "# set device\n",
    "device = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "print(device)\n",
    "\n",
    "# if run on a non-Apple chip, substitute the code above with the following:\n",
    "# device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NNClassifier to wrap PyTorch model and set hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier = NeuralNetClassifier(model,\n",
    "                                 max_epochs=80,\n",
    "                                 criterion=torch.nn.CrossEntropyLoss,\n",
    "                                 optimizer=torch.optim.Adam,\n",
    "                                 lr=0.0001,\n",
    "                                 train_split=None,\n",
    "                                 verbose=1,\n",
    "                                 device=device\n",
    "                                 )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mc_dropout_bald(classifier: BaseEstimator, X: modALinput, n_instances: int = 1,\n",
    "                    random_tie_break: bool = False, dropout_layer_indexes: list = [],\n",
    "                    num_cycles: int = 50, sample_per_forward_pass: int = 1000,\n",
    "                    logits_adaptor: Callable[[\n",
    "                        torch.tensor, modALinput], torch.tensor] = default_logits_adaptor,\n",
    "                    **mc_dropout_kwargs,) -> np.ndarray:\n",
    "    \"\"\"\n",
    "        Mc-Dropout bald query strategy. Returns the indexes of the instances with the largest BALD \n",
    "        (Bayesian Active Learning by Disagreement) score calculated through the dropout cycles\n",
    "        and the corresponding bald score. \n",
    "\n",
    "        Based on the work of: \n",
    "            Deep Bayesian Active Learning with Image Data.\n",
    "            (Yarin Gal, Riashat Islam, and Zoubin Ghahramani. 2017.)\n",
    "            Dropout as a Bayesian Approximation: Representing Model Uncer- tainty in Deep Learning.\n",
    "            (Yarin Gal and Zoubin Ghahramani. 2016.)\n",
    "            Bayesian Active Learning for Classification and Preference Learning.\n",
    "            (NeilHoulsby,FerencHusza ́r,ZoubinGhahramani,andMa ́te ́Lengyel. 2011.) \n",
    "\n",
    "        Args:\n",
    "            classifier: The classifier for which the labels are to be queried.\n",
    "            X: The pool of samples to query from.\n",
    "            n_instances: Number of samples to be queried.\n",
    "            random_tie_break: If True, shuffles utility scores to randomize the order. This\n",
    "                can be used to break the tie when the highest utility score is not unique.\n",
    "            dropout_layer_indexes: Indexes of the dropout layers which should be activated\n",
    "                Choose indices from : list(torch_model.modules())\n",
    "            num_cycles: Number of forward passes with activated dropout\n",
    "            sample_per_forward_pass: max. sample number for each forward pass. \n",
    "                The allocated RAM does mainly depend on this.\n",
    "                Small number --> small RAM allocation\n",
    "            logits_adaptor: Callable which can be used to adapt the output of a forward pass \n",
    "                to the required vector format for the vectorised metric functions \n",
    "            **uncertainty_measure_kwargs: Keyword arguments to be passed for the uncertainty\n",
    "                measure function.\n",
    "\n",
    "        Returns:\n",
    "            The indices of the instances from X chosen to be labelled;\n",
    "            The mc-dropout metric of the chosen instances; \n",
    "    \"\"\"\n",
    "    predictions = get_predictions(\n",
    "        classifier, X, dropout_layer_indexes, num_cycles, sample_per_forward_pass, logits_adaptor)\n",
    "    # calculate BALD (Bayesian active learning divergence))\n",
    "\n",
    "    bald_scores = _bald_divergence(predictions)\n",
    "\n",
    "    if not random_tie_break:\n",
    "        return multi_argmax(bald_scores, n_instances=n_instances)\n",
    "\n",
    "    return shuffled_argmax(bald_scores, n_instances=n_instances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_predictions(classifier: BaseEstimator, X: modALinput, dropout_layer_indexes: list = [],\n",
    "                    num_predictions: int = 50, sample_per_forward_pass: int = 1000,\n",
    "                    logits_adaptor: Callable[[torch.tensor, modALinput], torch.tensor] = default_logits_adaptor):\n",
    "    \"\"\"\n",
    "        Runs num_predictions times the prediction of the classifier on the input X \n",
    "        and puts the predictions in a list.\n",
    "\n",
    "        Args:\n",
    "            classifier: The classifier for which the labels are to be queried.\n",
    "            X: The pool of samples to query from.\n",
    "            dropout_layer_indexes: Indexes of the dropout layers which should be activated\n",
    "                Choose indices from : list(torch_model.modules())\n",
    "            num_predictions: Number of predictions which should be made\n",
    "            sample_per_forward_pass: max. sample number for each forward pass. \n",
    "                The allocated RAM does mainly depend on this.\n",
    "                Small number --> small RAM allocation\n",
    "            logits_adaptor: Callable which can be used to adapt the output of a forward pass \n",
    "                to the required vector format for the vectorised metric functions \n",
    "        Return: \n",
    "            prediction: list with all predictions\n",
    "    \"\"\"\n",
    "\n",
    "    assert num_predictions > 0, 'num_predictions must be larger than zero'\n",
    "    assert sample_per_forward_pass > 0, 'sample_per_forward_pass must be larger than zero'\n",
    "\n",
    "    predictions = []\n",
    "    # set dropout layers to train mode\n",
    "    set_dropout_mode(classifier.estimator.module_,\n",
    "                     dropout_layer_indexes, train_mode=True)\n",
    "\n",
    "    split_args = []\n",
    "\n",
    "    if isinstance(X, Mapping):  # check for dict\n",
    "        for k, v in X.items():\n",
    "\n",
    "            v.detach()\n",
    "            split_v = torch.split(v, sample_per_forward_pass)\n",
    "            # create sub-dictionary split for each forward pass with same keys&values\n",
    "            for split_idx, split in enumerate(split_v):\n",
    "                if len(split_args) <= split_idx:\n",
    "                    split_args.append({})\n",
    "                split_args[split_idx][k] = split\n",
    "\n",
    "    elif torch.is_tensor(X):  # check for tensor\n",
    "        X.detach()\n",
    "        split_args = torch.split(X, sample_per_forward_pass)\n",
    "    else:\n",
    "        raise RuntimeError(\n",
    "            \"Error in model data type, only dict or tensors supported\")\n",
    "\n",
    "    for i in range(num_predictions):\n",
    "\n",
    "        probas = []\n",
    "\n",
    "        for samples in split_args:\n",
    "            # call Skorch infer function to perform model forward pass\n",
    "            # In comparison to: predict(), predict_proba() the infer()\n",
    "            # does not change train/eval mode of other layers\n",
    "            with torch.no_grad():\n",
    "                logits = classifier.estimator.infer(samples)\n",
    "                prediction = logits_adaptor(logits, samples)\n",
    "                mask = ~prediction.isnan()\n",
    "                prediction[mask] = prediction[mask].softmax(-1)\n",
    "                probas.append(prediction)\n",
    "\n",
    "        probas = torch.cat(probas)\n",
    "        \n",
    "        # mps device wouldn't let me convert a pytorch tensor into a numpy array, because numpy operations on tensors in mps are not supported\n",
    "        # for this reason I moved the tensor to the CPU before converting it to a numpy array\n",
    "        if isinstance(probas, torch.Tensor):\n",
    "            probas = probas.cpu().detach().numpy()\n",
    "        else:\n",
    "            raise TypeError(\"probas must be a torch.Tensor\")\n",
    "\n",
    "        # Append the numpy array to predictions\n",
    "        predictions.append(probas)\n",
    "\n",
    "    # set dropout layers to eval\n",
    "    set_dropout_mode(classifier.estimator.module_,\n",
    "                     dropout_layer_indexes, train_mode=False)\n",
    "\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Active Learning Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize ActiveLearner (Pass the skorch wrapped PyTorch model and select query strategy)\n",
    "# query strategies:\n",
    "    # entropy_sampling\n",
    "    # margin_sampling\n",
    "    # mc_dropout_bald\n",
    "\n",
    "learner = DeepActiveLearner(\n",
    "    estimator=classifier,\n",
    "    query_strategy=mc_dropout_bald\n",
    ")\n",
    "\n",
    "# perform initial training\n",
    "model.train()\n",
    "learner.teach(X_initial, y_initial)\n",
    "\n",
    "print(\"Score from sklearn: {}\".format(learner.score(X_pool, y_pool)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose number of queries\n",
    "n_queries = 10\n",
    "\n",
    "X_teach = X_initial\n",
    "y_teach = y_initial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# confusion matrix before querying and after initial training\n",
    "before_train_predictions = learner.predict(X_teach)\n",
    "cm = confusion_matrix(y_teach, before_train_predictions)\n",
    "print(f\"Before training - Confusion Matrix\")\n",
    "print(cm)\n",
    "\n",
    "# plotting the confusion matrix\n",
    "plt.figure(figsize=(10, 8))\n",
    "sns.heatmap(cm, annot=True, fmt=\"d\", cmap=\"Blues\",\n",
    "            xticklabels=['Class 1', 'Class 2', 'Class 3', 'Class 10'],\n",
    "            yticklabels=['Class 1', 'Class 2', 'Class 3', 'Class 10'])\n",
    "\n",
    "plt.xticks([0.5, 1.5, 2.5, 9.5], labels=['Class 1', 'Class 2', 'Class 3', 'Class 10'], rotation=90)\n",
    "plt.yticks([0.5, 1.5, 2.5, 9.5], labels=['Class 1', 'Class 2', 'Class 3', 'Class 10'], rotation=0)\n",
    "\n",
    "plt.xlabel(\"Predicted\")\n",
    "plt.ylabel(\"True\")\n",
    "plt.title(f\"Before trainining - Confusion Matrix\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for run_number in range(10):\n",
    "    # initialize wandb tracking, set configuration\n",
    "    wandb.init(\n",
    "        project=\"MC_BALD_KMNIST\",\n",
    "        config={\n",
    "            \"model\": \"CNNmodel\",\n",
    "            \"dataset\": \"KMNIST\",\n",
    "            \"split\": \"48k/12k/10k\",\n",
    "            \"batch_size\": 60000,\n",
    "            \"loss\": \"CrossEntropyLoss\",\n",
    "            \"optimization\": \"Adam\",\n",
    "            \"learning_rate\": 0.0001,\n",
    "            \"dropout\": \"0.4/0.7\",\n",
    "            \"epochs\": 80,\n",
    "            \"n_queries\": 10,\n",
    "            \"initial\": 1000,\n",
    "            \"n_instances\": 100,\n",
    "            \"num_cycles\" : 50\n",
    "        },\n",
    "    )\n",
    "    wandb_step = 1\n",
    "        \n",
    "    for idx in range(n_queries):\n",
    "        print('Query no. %d' % (idx + 1))\n",
    "        \n",
    "        # adjust model.train()/ model.eval() according to the strategy and experimental setup:\n",
    "        # for non-MC dropout strategies, I set the model to model.eval() before querying\n",
    "        # BALD MC dropout will automatically turn on training mode before calculating uncertainties\n",
    "        # model.train will keep dropout layers on during inference, whereas model.eval will deactivate them during inference\n",
    "        # mc_dropout_bald automatically sets model.train() before selecting instances to label\n",
    "        \n",
    "        query_idx, metric_values = learner.query(X_pool, n_instances=100, sample_per_forward_pass=250, num_cycles=50)\n",
    "        \n",
    "        # NB: for non-mc_dropout_bald strategies, remove \"sample_per_forward_pass, num_cycles\"\n",
    "        # n_instances = n of samples to be queried\n",
    "        # all dropout layers will be activated on default. \n",
    "        # sample_per_forward_pass = max. sample number for each forward pass\n",
    "        # num_cycles are the number of dropout forward passes that should be performed\n",
    "        \n",
    "        # Add queried instances\n",
    "        X_teach = torch.cat((X_teach, X_pool[query_idx]))\n",
    "        y_teach = torch.cat((y_teach, y_pool[query_idx]))\n",
    "        \n",
    "        # set training mode when updating the model with labelled instances\n",
    "        model.train()\n",
    "        learner.teach(X_teach, y_teach)\n",
    "\n",
    "        # remove queried instance from pool\n",
    "        X_pool = np.delete(X_pool, query_idx, axis=0)\n",
    "        y_pool = np.delete(y_pool, query_idx, axis=0)\n",
    "        \n",
    "        # Evaluate and log training metrics\n",
    "        model.eval()\n",
    "        train_predictions = learner.predict(X_teach)\n",
    "        train_accuracy = accuracy_score(y_teach, train_predictions)*100\n",
    "        \n",
    "        train_f1 = f1_score(y_teach, train_predictions, average='weighted')*100\n",
    "        train_precision = precision_score(y_teach, train_predictions, average='weighted')*100\n",
    "        train_recall = recall_score(y_teach, train_predictions, average='weighted')*100\n",
    "        \n",
    "        print(\"Train Accuracy: {}\".format(train_accuracy))\n",
    "        print(\"Train F1 Score: {:.4f}\".format(train_f1))\n",
    "        print(\"Train Precision: {:.4f}\".format(train_precision))\n",
    "        print(\"Train Recall: {:.4f}\".format(train_recall))\n",
    "        \n",
    "        wandb.log({\"Train Accuracy\": train_accuracy, \"Train F1 Score\": train_f1, \"Train Precision\": train_precision, \"Train Recall\": train_recall}, step=wandb_step)\n",
    "        \n",
    "        # Evaluate and log validation metrics\n",
    "        model.eval()\n",
    "        val_predictions = learner.predict(X_val)\n",
    "        val_accuracy = accuracy_score(y_val, val_predictions)*100\n",
    "        val_f1 = f1_score(y_val, val_predictions, average='weighted')*100\n",
    "        val_precision = precision_score(y_val, val_predictions, average='weighted')*100\n",
    "        val_recall = recall_score(y_val, val_predictions, average='weighted')*100\n",
    "\n",
    "        print(\"Val Accuracy: {}\".format(val_accuracy))\n",
    "        print(\"Val F1 Score: {:.4f}\".format(val_f1))\n",
    "        print(\"Val Precision: {:.4f}\".format(val_precision))\n",
    "        print(\"Val Recall: {:.4f}\".format(val_recall))\n",
    "\n",
    "        wandb.log({\"Val Accuracy\": val_accuracy, \"Val F1 Score\": val_f1, \"Val Precision\": val_precision, \"Val Recall\": val_recall}, step=wandb_step)\n",
    "        \n",
    "        wandb_step += 1\n",
    "    \n",
    "    # evaluate on test dataset with dropout off, as they did in Gal et al. (2017)\n",
    "    model.eval()\n",
    "    # Evaluate and log test metrics after the loop\n",
    "    test_predictions = learner.predict(X_test)\n",
    "    test_accuracy = accuracy_score(y_test, test_predictions)*100\n",
    "    test_f1 = f1_score(y_test, test_predictions, average='weighted')*100\n",
    "    test_precision = precision_score(y_test, test_predictions, average='weighted')*100\n",
    "    test_recall = recall_score(y_test, test_predictions, average='weighted')*100\n",
    "\n",
    "    print(\"Test Accuracy: {}\".format(test_accuracy))\n",
    "    print(\"Test F1 Score: {:.4f}\".format(test_f1))\n",
    "    print(\"Test Precision: {:.4f}\".format(test_precision))\n",
    "    print(\"Test Recall: {:.4f}\".format(test_recall))\n",
    "\n",
    "    wandb.log({\"Test Accuracy\": test_accuracy, \"Test F1 Score\": test_f1, \"Test Precision\": test_precision, \"Test Recall\": test_recall})\n",
    "\n",
    "    wandb.finish()\n",
    "\n",
    "    # After running the AL loop, calculate confusion matrix and print report on test set\n",
    "    \n",
    "    cm = confusion_matrix(y_test, test_predictions)\n",
    "    print(f\"Confusion Matrix - Run {run_number + 1}:\")\n",
    "    print(cm)\n",
    "\n",
    "    # Plotting the confusion matrix\n",
    "    plt.figure(figsize=(10, 8))\n",
    "    sns.heatmap(cm, annot=True, fmt=\"d\", cmap=\"Blues\",\n",
    "                xticklabels=['Class 1', 'Class 2', 'Class 3', 'Class 10'],\n",
    "                yticklabels=['Class 1', 'Class 2', 'Class 3', 'Class 10'])\n",
    "\n",
    "    plt.xticks([0.5, 1.5, 2.5, 9.5], labels=['Class 1', 'Class 2', 'Class 3', 'Class 10'], rotation=90)\n",
    "    plt.yticks([0.5, 1.5, 2.5, 9.5], labels=['Class 1', 'Class 2', 'Class 3', 'Class 10'], rotation=0)\n",
    "\n",
    "    plt.xlabel(\"Predicted\")\n",
    "    plt.ylabel(\"True\")\n",
    "    plt.title(f\"Run {run_number + 1} - Confusion Matrix\")\n",
    "    plt.show()\n",
    "\n",
    "    # Print classification report\n",
    "    print(f\"Classification Report - Run {run_number + 1}:\")\n",
    "    print(classification_report(y_test, test_predictions))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlproj",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}