CIFAR10_ResNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ce3002-866a-47f2-a1f7-ae220354568a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A lot of code taken from PyTorch Lightning Tutorial\n",
    "import os\n",
    "import urllib.request\n",
    "from types import SimpleNamespace\n",
    "from urllib.error import HTTPError\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "import seaborn as sns\n",
    "import tabulate\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "import torchvision\n",
    "\n",
    "%matplotlib inline\n",
    "from IPython.display import HTML, display, set_matplotlib_formats\n",
    "from PIL import Image\n",
    "from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n",
    "from random import choices\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "from utils import weighted_mean, weighted_sd, read_results\n",
    "from torch_utils import d_LSE, BI_LSE, BI_thresh_acc_DE, BI_thresh_acc_TTA, Conf_thresh_acc, BI_thresh_acc_MCD, BI_thresh_DE, Conf_thresh\n",
    "\n",
    "import torch.nn.functional as F\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "set_matplotlib_formats(\"svg\", \"pdf\")  # For export\n",
    "matplotlib.rcParams[\"lines.linewidth\"] = 2.0\n",
    "sns.reset_orig()\n",
    "\n",
    "os.makedirs('plots', exist_ok=True)\n",
    "\n",
    "ds_name = 'Cifar10-C'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f93758ed-2138-4cb9-8812-dca8aba20322",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the folder where the datasets are/should be downloaded (e.g. CIFAR10)\n",
    "DATASET_PATH = os.environ.get(\"PATH_DATASETS\", \"../data/\")\n",
    "# Path to the folder where the pretrained models are saved\n",
    "CHECKPOINT_PATH = os.environ.get(\"PATH_CHECKPOINT\", \"../saved_models/ConvNets\")\n",
    "os.makedirs(CHECKPOINT_PATH, exist_ok=True)\n",
    "\n",
    "# Function for setting the seed\n",
    "pl.seed_everything(42)\n",
    "\n",
    "# Ensure that all operations are deterministic on GPU (if used) for reproducibility\n",
    "torch.backends.cudnn.determinstic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "\n",
    "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2cf08f7-e0fd-429c-8690-d9d9ba2fefa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = CIFAR10(root=DATASET_PATH, train=True, download=True)\n",
    "DATA_MEANS = (train_dataset.data / 255.0).mean(axis=(0, 1, 2))\n",
    "DATA_STD = (train_dataset.data / 255.0).std(axis=(0, 1, 2))\n",
    "print(\"Data mean\", DATA_MEANS)\n",
    "print(\"Data std\", DATA_STD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c08b511-f9ff-48a1-9b5b-f82b5b5fa095",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(DATA_MEANS, DATA_STD)])\n",
    "# For training, we add some augmentation. Networks are too powerful and would overfit.\n",
    "train_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.RandomResizedCrop((32, 32), scale=(0.8, 1.0), ratio=(0.9, 1.1)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(DATA_MEANS, DATA_STD),\n",
    "    ]\n",
    ")\n",
    "# Loading the training dataset. We need to split it into a training and validation part\n",
    "# We need to do a little trick because the validation set should not use the augmentation.\n",
    "train_dataset = CIFAR10(root=DATASET_PATH, train=True, transform=train_transform, download=True)\n",
    "val_dataset = CIFAR10(root=DATASET_PATH, train=True, transform=test_transform, download=True)\n",
    "pl.seed_everything(42)\n",
    "train_set, _ = torch.utils.data.random_split(train_dataset, [45000, 5000])\n",
    "pl.seed_everything(42)\n",
    "_, val_set = torch.utils.data.random_split(val_dataset, [45000, 5000])\n",
    "\n",
    "# Loading the test set\n",
    "test_set = CIFAR10(root=DATASET_PATH, train=False, transform=test_transform, download=True)\n",
    "\n",
    "# We define a set of data loaders that we can use for various purposes later.\n",
    "train_loader = data.DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True, pin_memory=True, num_workers=4)\n",
    "val_loader = data.DataLoader(val_set, batch_size=128, shuffle=False, drop_last=False, num_workers=4)\n",
    "test_loader = data.DataLoader(test_set, batch_size=128, shuffle=False, drop_last=False, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ad5cc59-1962-47e0-8f85-ad74ad2c4e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CIFARModule(pl.LightningModule):\n",
    "    def __init__(self, model_name, model_hparams, optimizer_name, optimizer_hparams):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            model_name - Name of the model/CNN to run. Used for creating the model (see function below)\n",
    "            model_hparams - Hyperparameters for the model, as dictionary.\n",
    "            optimizer_name - Name of the optimizer to use. Currently supported: Adam, SGD\n",
    "            optimizer_hparams - Hyperparameters for the optimizer, as dictionary. This includes learning rate, weight decay, etc.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        # Exports the hyperparameters to a YAML file, and create \"self.hparams\" namespace\n",
    "        self.save_hyperparameters()\n",
    "        # Create model\n",
    "        self.model = create_model(model_name, model_hparams)\n",
    "        # Create loss module\n",
    "        self.loss_module = nn.CrossEntropyLoss()\n",
    "        # Example input for visualizing the graph in Tensorboard\n",
    "        self.example_input_array = torch.zeros((1, 3, 32, 32), dtype=torch.float32)\n",
    "\n",
    "    def forward(self, imgs):\n",
    "        # Forward function that is run when visualizing the graph\n",
    "        return self.model(imgs)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        # We will support Adam or SGD as optimizers.\n",
    "        if self.hparams.optimizer_name == \"Adam\":\n",
    "            # AdamW is Adam with a correct implementation of weight decay (see here\n",
    "            # for details: https://arxiv.org/pdf/1711.05101.pdf)\n",
    "            optimizer = optim.AdamW(self.parameters(), **self.hparams.optimizer_hparams)\n",
    "        elif self.hparams.optimizer_name == \"SGD\":\n",
    "            optimizer = optim.SGD(self.parameters(), **self.hparams.optimizer_hparams)\n",
    "        else:\n",
    "            assert False, f'Unknown optimizer: \"{self.hparams.optimizer_name}\"'\n",
    "\n",
    "        # We will reduce the learning rate by 0.1 after 100 and 150 epochs\n",
    "        scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1)\n",
    "        return [optimizer], [scheduler]\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        # \"batch\" is the output of the training data loader.\n",
    "        imgs, labels = batch\n",
    "        preds = self.model(imgs)\n",
    "        loss = self.loss_module(preds, labels)\n",
    "        acc = (preds.argmax(dim=-1) == labels).float().mean()\n",
    "\n",
    "        # Logs the accuracy per epoch to tensorboard (weighted average over batches)\n",
    "        self.log(\"train_acc\", acc, on_step=False, on_epoch=True)\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss  # Return tensor to call \".backward\" on\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        imgs, labels = batch\n",
    "        preds = self.model(imgs).argmax(dim=-1)\n",
    "        acc = (labels == preds).float().mean()\n",
    "        # By default logs it per epoch (weighted average over batches)\n",
    "        self.log(\"val_acc\", acc)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        imgs, labels = batch\n",
    "        preds = self.model(imgs).argmax(dim=-1)\n",
    "        acc = (labels == preds).float().mean()\n",
    "        # By default logs it per epoch (weighted average over batches), and returns it afterwards\n",
    "        self.log(\"test_acc\", acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ecf3bb7-7764-4b47-bdf2-344309dc9263",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict = {}\n",
    "\n",
    "\n",
    "def create_model(model_name, model_hparams):\n",
    "    if model_name in model_dict:\n",
    "        return model_dict[model_name](**model_hparams)\n",
    "    else:\n",
    "        assert False, f'Unknown model name \"{model_name}\". Available models are: {str(model_dict.keys())}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bd3aec2-1324-4da6-9775-5dfd8ff0ac0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "act_fn_by_name = {\"tanh\": nn.Tanh, \"relu\": nn.ReLU, \"leakyrelu\": nn.LeakyReLU, \"gelu\": nn.GELU}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf0da70-d582-4ca7-b246-b12a888b9506",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model_name, save_name=None, seed=42, max_epochs=180, verbose=True, **kwargs):\n",
    "    \"\"\"\n",
    "    Inputs:\n",
    "        model_name - Name of the model you want to run. Is used to look up the class in \"model_dict\"\n",
    "        save_name (optional) - If specified, this name will be used for creating the checkpoint and logging directory.\n",
    "    \"\"\"\n",
    "    if save_name is None:\n",
    "        save_name = model_name + '_seed{}'.format(seed)\n",
    "\n",
    "    # Create a PyTorch Lightning trainer with the generation callback\n",
    "    trainer = pl.Trainer(\n",
    "        default_root_dir=os.path.join(CHECKPOINT_PATH, save_name),  # Where to save models\n",
    "        accelerator='gpu',\n",
    "        devices=1,\n",
    "        # How many epochs to train for if no patience is set\n",
    "        max_epochs=max_epochs,\n",
    "        callbacks=[\n",
    "            ModelCheckpoint(\n",
    "                save_weights_only=True, mode=\"max\", monitor=\"val_acc\"\n",
    "            ),  # Save the best checkpoint based on the maximum val_acc recorded. Saves only weights and not optimizer\n",
    "            LearningRateMonitor(\"epoch\")\n",
    "        ],  # Log learning rate every epoch\n",
    "        enable_progress_bar=verbose,\n",
    "        logger=verbose,\n",
    "    )  # In case your notebook crashes due to the progress bar, consider increasing the refresh rate\n",
    "    if verbose:\n",
    "        trainer.logger._log_graph = True  # If True, we plot the computation graph in tensorboard\n",
    "        trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need\n",
    "\n",
    "    # Check whether pretrained model exists. If yes, load it and skip training\n",
    "    pretrained_filename = os.path.join(CHECKPOINT_PATH, save_name + \".ckpt\")\n",
    "    if os.path.isfile(pretrained_filename):\n",
    "        print(f\"Found pretrained model at {pretrained_filename}, loading...\")\n",
    "        # Automatically loads the model with the saved hyperparameters\n",
    "        model = CIFARModule.load_from_checkpoint(pretrained_filename)\n",
    "    else:\n",
    "        pl.seed_everything(seed)\n",
    "        model = CIFARModule(model_name=model_name, **kwargs)\n",
    "        trainer.fit(model, train_loader, val_loader)\n",
    "        model = CIFARModule.load_from_checkpoint(\n",
    "            trainer.checkpoint_callback.best_model_path\n",
    "        )  # Load best checkpoint after training\n",
    "        trainer.save_checkpoint(pretrained_filename)\n",
    "\n",
    "    # Test best model on validation and test set\n",
    "    val_result = trainer.test(model, dataloaders=val_loader, verbose=False)\n",
    "    test_result = trainer.test(model, dataloaders=test_loader, verbose=False)\n",
    "    result = {\"test\": test_result[0][\"test_acc\"], \"val\": val_result[0][\"test_acc\"]}\n",
    "\n",
    "    return model, result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25932b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_BS_model(model_name, save_name=None, seed=42, max_epochs=180, verbose=True, **kwargs):\n",
    "    \"\"\"\n",
    "    Inputs:\n",
    "        model_name - Name of the model you want to run. Is used to look up the class in \"model_dict\"\n",
    "        save_name (optional) - If specified, this name will be used for creating the checkpoint and logging directory.\n",
    "    \"\"\"\n",
    "    if save_name is None:\n",
    "        save_name = model_name + '_BS_seed{}'.format(seed)\n",
    "\n",
    "    # Create a PyTorch Lightning trainer with the generation callback\n",
    "    trainer = pl.Trainer(\n",
    "        default_root_dir=os.path.join(CHECKPOINT_PATH, save_name),  # Where to save models\n",
    "        accelerator='gpu',\n",
    "        devices=1,\n",
    "        # How many epochs to train for if no patience is set\n",
    "        max_epochs=max_epochs,\n",
    "        callbacks=[\n",
    "            ModelCheckpoint(\n",
    "                save_weights_only=True, mode=\"max\", monitor=\"val_acc\"\n",
    "            ),  # Save the best checkpoint based on the maximum val_acc recorded. Saves only weights and not optimizer\n",
    "            #LearningRateMonitor(\"epoch\")\n",
    "        ],  # Log learning rate every epoch\n",
    "        enable_progress_bar=verbose,\n",
    "        logger=verbose,\n",
    "    )  # In case your notebook crashes due to the progress bar, consider increasing the refresh rate\n",
    "    if verbose:\n",
    "        trainer.logger._log_graph = True  # If True, we plot the computation graph in tensorboard\n",
    "        trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need\n",
    "\n",
    "    # Check whether pretrained model exists. If yes, load it and skip training\n",
    "    pretrained_filename = os.path.join(CHECKPOINT_PATH, save_name + \".ckpt\")\n",
    "    if os.path.isfile(pretrained_filename):\n",
    "        print(f\"Found pretrained model at {pretrained_filename}, loading...\")\n",
    "        # Automatically loads the model with the saved hyperparameters\n",
    "        model = CIFARModule.load_from_checkpoint(pretrained_filename)\n",
    "    else:\n",
    "        pl.seed_everything(seed)\n",
    "        BS_indices = choices(range(len(train_set)), k=len(train_set))\n",
    "        BS_set = data.Subset(train_set, BS_indices)\n",
    "        train_loader = data.DataLoader(\n",
    "            BS_set, batch_size=128, shuffle=True, drop_last=True, pin_memory=True, num_workers=4)\n",
    "        model = CIFARModule(model_name=model_name, **kwargs)\n",
    "        trainer.fit(model, train_loader, val_loader)\n",
    "        model = CIFARModule.load_from_checkpoint(\n",
    "            trainer.checkpoint_callback.best_model_path\n",
    "        )  # Load best checkpoint after training\n",
    "        trainer.save_checkpoint(pretrained_filename)\n",
    "\n",
    "    # Test best model on validation and test set\n",
    "    val_result = trainer.test(model, dataloaders=val_loader, verbose=False)\n",
    "    test_result = trainer.test(model, dataloaders=test_loader, verbose=False)\n",
    "    result = {\"test\": test_result[0][\"test_acc\"], \"val\": val_result[0][\"test_acc\"]}\n",
    "\n",
    "    return model, result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62e80df2-3f99-4abf-bf3d-ed35697a6912",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNetBlock(nn.Module):\n",
    "    def __init__(self, c_in, act_fn, subsample=False, c_out=-1):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            c_in - Number of input features\n",
    "            act_fn - Activation class constructor (e.g. nn.ReLU)\n",
    "            subsample - If True, we want to apply a stride inside the block and reduce the output shape by 2 in height and width\n",
    "            c_out - Number of output features. Note that this is only relevant if subsample is True, as otherwise, c_out = c_in\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        if not subsample:\n",
    "            c_out = c_in\n",
    "\n",
    "        # Network representing F\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                c_in, c_out, kernel_size=3, padding=1, stride=1 if not subsample else 2, bias=False\n",
    "            ),  # No bias needed as the Batch Norm handles it\n",
    "            nn.BatchNorm2d(c_out),\n",
    "            act_fn(),\n",
    "            nn.Conv2d(c_out, c_out, kernel_size=3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(c_out),\n",
    "        )\n",
    "\n",
    "        # 1x1 convolution with stride 2 means we take the upper left value, and transform it to new output size\n",
    "        self.downsample = nn.Conv2d(c_in, c_out, kernel_size=1, stride=2) if subsample else None\n",
    "        self.act_fn = act_fn()\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.net(x)\n",
    "        if self.downsample is not None:\n",
    "            x = self.downsample(x)\n",
    "        out = z + x\n",
    "        out = self.act_fn(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b10f857f-648c-4216-8ca6-c3e4c0199050",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PreActResNetBlock(nn.Module):\n",
    "    def __init__(self, c_in, act_fn, subsample=False, c_out=-1):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            c_in - Number of input features\n",
    "            act_fn - Activation class constructor (e.g. nn.ReLU)\n",
    "            subsample - If True, we want to apply a stride inside the block and reduce the output shape by 2 in height and width\n",
    "            c_out - Number of output features. Note that this is only relevant if subsample is True, as otherwise, c_out = c_in\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        if not subsample:\n",
    "            c_out = c_in\n",
    "\n",
    "        # Network representing F\n",
    "        self.net = nn.Sequential(\n",
    "            nn.BatchNorm2d(c_in),\n",
    "            act_fn(),\n",
    "            nn.Conv2d(c_in, c_out, kernel_size=3, padding=1, stride=1 if not subsample else 2, bias=False),\n",
    "            nn.BatchNorm2d(c_out),\n",
    "            act_fn(),\n",
    "            nn.Conv2d(c_out, c_out, kernel_size=3, padding=1, bias=False),\n",
    "        )\n",
    "\n",
    "        # 1x1 convolution needs to apply non-linearity as well as not done on skip connection\n",
    "        self.downsample = (\n",
    "            nn.Sequential(nn.BatchNorm2d(c_in), act_fn(), nn.Conv2d(c_in, c_out, kernel_size=1, stride=2, bias=False))\n",
    "            if subsample\n",
    "            else None\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.net(x)\n",
    "        if self.downsample is not None:\n",
    "            x = self.downsample(x)\n",
    "        out = z + x\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d62d15e-c762-4666-a6e4-9963ce3088a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_blocks_by_name = {\"ResNetBlock\": ResNetBlock, \"PreActResNetBlock\": PreActResNetBlock}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "445d3a00-ee36-4b27-acda-d2dfa50e8e8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_classes=10,\n",
    "        num_blocks=[3, 3, 3],\n",
    "        c_hidden=[16, 32, 64],\n",
    "        act_fn_name=\"relu\",\n",
    "        block_name=\"ResNetBlock\",\n",
    "        **kwargs,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Inputs:\n",
    "            num_classes - Number of classification outputs (10 for CIFAR10)\n",
    "            num_blocks - List with the number of ResNet blocks to use. The first block of each group uses downsampling, except the first.\n",
    "            c_hidden - List with the hidden dimensionalities in the different blocks. Usually multiplied by 2 the deeper we go.\n",
    "            act_fn_name - Name of the activation function to use, looked up in \"act_fn_by_name\"\n",
    "            block_name - Name of the ResNet block, looked up in \"resnet_blocks_by_name\"\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        assert block_name in resnet_blocks_by_name\n",
    "        self.hparams = SimpleNamespace(\n",
    "            num_classes=num_classes,\n",
    "            c_hidden=c_hidden,\n",
    "            num_blocks=num_blocks,\n",
    "            act_fn_name=act_fn_name,\n",
    "            act_fn=act_fn_by_name[act_fn_name],\n",
    "            block_class=resnet_blocks_by_name[block_name],\n",
    "        )\n",
    "        self._create_network()\n",
    "        self._init_params()\n",
    "\n",
    "    def _create_network(self):\n",
    "        c_hidden = self.hparams.c_hidden\n",
    "\n",
    "        # A first convolution on the original image to scale up the channel size\n",
    "        if self.hparams.block_class == PreActResNetBlock:  # => Don't apply non-linearity on output\n",
    "            self.input_net = nn.Sequential(nn.Conv2d(3, c_hidden[0], kernel_size=3, padding=1, bias=False))\n",
    "        else:\n",
    "            self.input_net = nn.Sequential(\n",
    "                nn.Conv2d(3, c_hidden[0], kernel_size=3, padding=1, bias=False),\n",
    "                nn.BatchNorm2d(c_hidden[0]),\n",
    "                self.hparams.act_fn(),\n",
    "            )\n",
    "\n",
    "        # Creating the ResNet blocks\n",
    "        blocks = []\n",
    "        for block_idx, block_count in enumerate(self.hparams.num_blocks):\n",
    "            for bc in range(block_count):\n",
    "                # Subsample the first block of each group, except the very first one.\n",
    "                subsample = bc == 0 and block_idx > 0\n",
    "                blocks.append(\n",
    "                    self.hparams.block_class(\n",
    "                        c_in=c_hidden[block_idx if not subsample else (block_idx - 1)],\n",
    "                        act_fn=self.hparams.act_fn,\n",
    "                        subsample=subsample,\n",
    "                        c_out=c_hidden[block_idx],\n",
    "                    )\n",
    "                )\n",
    "        self.blocks = nn.Sequential(*blocks)\n",
    "\n",
    "        # Mapping to classification output\n",
    "        self.output_net = nn.Sequential(\n",
    "            nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(c_hidden[-1], self.hparams.num_classes)\n",
    "        )\n",
    "\n",
    "    def _init_params(self):\n",
    "        # Based on our discussion in Tutorial 4, we should initialize the convolutions according to the activation function\n",
    "        # Fan-out focuses on the gradient distribution, and is commonly used in ResNets\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode=\"fan_out\", nonlinearity=self.hparams.act_fn_name)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.input_net(x)\n",
    "        x = self.blocks(x)\n",
    "        x = self.output_net(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60b2cf8a-326e-4392-9bd9-602d0f95ded3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict[\"ResNet\"] = ResNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574e8bd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_ens = 10\n",
    "max_epochs = 180\n",
    "\n",
    "ensemble = {'models': [], 'results': [], 'BS_models': [], 'BS_results': []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1df65d30",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Bootstrapping (not featured in the paper - did not improve the results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e40632e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for seed in range(n_ens):\n",
    "#     resnet_model, resnet_results = train_BS_model(\n",
    "#         model_name=\"ResNet\",\n",
    "#         model_hparams={\"num_classes\": 10, \"c_hidden\": [16, 32, 64], \"num_blocks\": [3, 3, 3], \"act_fn_name\": \"relu\"},\n",
    "#         optimizer_name=\"SGD\",\n",
    "#         optimizer_hparams={\"lr\": 0.1, \"momentum\": 0.9, \"weight_decay\": 1e-4},\n",
    "#         seed=seed,\n",
    "#         max_epochs=max_epochs,\n",
    "#         verbose=False\n",
    "#     )\n",
    "#     ensemble['BS_models'].append(resnet_model.to(device='cpu'))\n",
    "#     ensemble['BS_results'].append(resnet_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b0fb1ea-befb-4e56-801d-9df20c4b04f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Deep ensembles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e1e36f-97d9-444a-9637-56e073c58e73",
   "metadata": {},
   "outputs": [],
   "source": [
    "for seed in range(n_ens):\n",
    "    resnet_model, resnet_results = train_model(\n",
    "        model_name=\"ResNet\",\n",
    "        model_hparams={\"num_classes\": 10, \"c_hidden\": [16, 32, 64], \"num_blocks\": [3, 3, 3], \"act_fn_name\": \"relu\"},\n",
    "        optimizer_name=\"SGD\",\n",
    "        optimizer_hparams={\"lr\": 0.1, \"momentum\": 0.9, \"weight_decay\": 1e-4},\n",
    "        seed=seed,\n",
    "        max_epochs=max_epochs,\n",
    "        verbose=False\n",
    "    )\n",
    "    ensemble['models'].append(resnet_model.to(device='cpu'))\n",
    "    ensemble['results'].append(resnet_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843753a5-8778-4376-b393-f16986ecac5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test and validation accuracy of each ensemble member\n",
    "ensemble['results']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa8d0d2c-d737-400b-a211-9b4fb68b70a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CIFAR10-C\n",
    "\n",
    "def imshow(loader):\n",
    "    dataiter = iter(loader)\n",
    "    images, labels = dataiter.next()\n",
    "    img = torchvision.utils.make_grid(images)\n",
    "    img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "\n",
    "def get_loader(corruption, batch_size=128, severity=5):\n",
    "    xs = transforms.Normalize(DATA_MEANS, DATA_STD)(torch.from_numpy(np.load('../data/CIFAR-10-C/{}.npy'.format(corruption)) / 255.).float().transpose(1, 2).transpose(1, 3))\n",
    "    ys = torch.from_numpy(np.load('../data/CIFAR-10-C/labels.npy'))\n",
    "    # only pick the last 10000 as these have the highest corruption level\n",
    "    #n_data = xs.shape[0]\n",
    "    xs = xs[(10000*(severity-1)):(10000*severity)]\n",
    "    ys = ys[(10000*(severity-1)):(10000*severity)]\n",
    "\n",
    "    ood_set = torch.utils.data.TensorDataset(xs, ys)\n",
    "    # We define a set of data loaders that we can use for various purposes later.\n",
    "    return torch.utils.data.DataLoader(ood_set, batch_size=batch_size, shuffle=True, drop_last=False, num_workers=4)\n",
    "\n",
    "corruptions = [\n",
    "    'None',\n",
    "    'gaussian_blur',\n",
    "    'brightness',\n",
    "    'speckle_noise',\n",
    "    'glass_blur',\n",
    "    'spatter',\n",
    "    'shot_noise',\n",
    "    'defocus_blur',\n",
    "    'elastic_transform',\n",
    "    'frost',\n",
    "    'saturate',\n",
    "    'snow',\n",
    "    'gaussian_noise',\n",
    "    'contrast',\n",
    "    'motion_blur',\n",
    "    'impulse_noise',\n",
    "    'pixelate',\n",
    "    'fog',\n",
    "    'jpeg_compression',\n",
    "    'zoom_blur',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb6e089-0304-43b3-a12a-580d9472b26e",
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow(get_loader('speckle_noise', batch_size=16))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaf4841d-9a04-4dd6-9bab-65b9cd2b14d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This cell takes ~1 hour (single RTX 5000)\n",
    "qs = np.arange(10, -1, -1) / 10\n",
    "# pick the first model in the ensemble to get the accuracy\n",
    "model = ensemble['models'][0]\n",
    "for target in tqdm(['NLL', 'Accuracy'], desc='Target'):\n",
    "    for type_ in tqdm(['DE', 'TTA', 'Conf'], desc='Type'):\n",
    "\n",
    "        if type_=='DE':\n",
    "            experiment = BI_thresh_DE(model, ensemble['models'], [val_loader], qs=qs, device='cuda:0')\n",
    "        elif type_=='BS':\n",
    "            experiment = BI_thresh_DE(model, ensemble['BS_models'], [val_loader], qs=qs, device='cuda:0')\n",
    "        elif type_=='TTA':\n",
    "            experiment = BI_thresh_acc_TTA(model, [val_loader], qs=qs, n_ens=20, device='cuda:0', img_size=32)\n",
    "        elif type_=='Conf':\n",
    "            experiment = Conf_thresh(model, [val_loader], qs=qs, device='cuda:0')\n",
    "        elif type_=='MCD':\n",
    "            experiment = BI_thresh_acc_MCD(model, [val_loader], qs=qs, n_ens=100, device='cuda:0')\n",
    "\n",
    "        for severity in tqdm(range(1, 6), leave=False, desc='severity'):\n",
    "            results = []\n",
    "            for corr in tqdm(corruptions, leave=False, desc='Corruptions'):\n",
    "                if corr == 'None':\n",
    "                    results_ = experiment.run([test_loader])\n",
    "                else:\n",
    "                    ood_loader = get_loader(corr, severity=severity)\n",
    "                    results_ = experiment.run([ood_loader])\n",
    "                results_['Corruption'] = corr\n",
    "                \n",
    "                os.makedirs('results/{}/{}/'.format(ds_name, target), exist_ok=True)\n",
    "                results_.to_pickle('results/{}/{}/{}_sev{}_corr{}.pkl'.format(ds_name, target, type_, severity, corr))\n",
    "                results.append(results_)\n",
    "\n",
    "            results = pd.concat(results)\n",
    "            results = results.reset_index(drop=True)\n",
    "            results.to_pickle('results/{}/{}/{}_sev{}_all.pkl'.format(ds_name, target, type_, severity))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1a98f7d-53d6-42dc-b3b0-dc179b27dc28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure 11b\n",
    "sns.set_palette([ '#1f77b4', '#d62728',])\n",
    "target = 'NLL'\n",
    "unc_type = 'DE'\n",
    "hue = 'Uncertainty'\n",
    "\n",
    "for severity in range(1, 6):\n",
    "    results_DE = read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)\n",
    "    #results_BS = read_results(unc_type='BS', ds_name=ds_name, severity=severity, target=target)\n",
    "    results_Conf = read_results(unc_type='Conf', ds_name=ds_name, severity=severity, target=target)\n",
    "    #results_combined = pd.concat([results_Conf, results_DE, results_BS])\n",
    "    results_combined = pd.concat([results_Conf, results_DE])\n",
    "    results_combined['Test set'] = results_combined['Corruption Severity'].replace({\n",
    "        '0': 'CIFAR10',\n",
    "        '1': 'CIFAR10-C (Severity 1)',\n",
    "        '2': 'CIFAR10-C (Severity 2)',\n",
    "        '3': 'CIFAR10-C (Severity 3)',\n",
    "        '4': 'CIFAR10-C (Severity 4)',\n",
    "        '5': 'CIFAR10-C (Highest severity)',\n",
    "    })\n",
    "\n",
    "    plt.rcParams.update({'font.size': 12, 'figure.dpi': 1000})\n",
    "    figure = plt.figure(figsize=(6, 4))\n",
    "    ax = plt.subplot(1, 1, 1)\n",
    "    plot = sns.lineplot(\n",
    "        data=results_combined, x='Validation set quantile', y=target, hue=hue, style='Test set', ax=ax, linewidth=3, \n",
    "        estimator=weighted_mean, errorbar=weighted_sd)\n",
    "    plot.invert_xaxis()\n",
    "    sns.despine()\n",
    "    plt.xlim(1.02, 0.4)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig('plots/resnet_{}_quantile_hue{}_sev{}_unc{}_{}.png'.format(ds_name, hue, severity, unc_type, target), bbox_inches='tight')\n",
    "    plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236f5d8c-2193-4ed8-92ef-08c1ab1ddf7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure 1 & 12c\n",
    "unc_type = 'DE'\n",
    "# Figure 12d\n",
    "#unc_type = 'TTA'\n",
    "\n",
    "sns.set_palette([ '#1f77b4', '#d62728',])\n",
    "target = 'Accuracy'\n",
    "#unc_type = 'TTA'\n",
    "hue = 'Uncertainty'\n",
    "\n",
    "for severity in range(1, 6):\n",
    "    results_DE = read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)\n",
    "    #results_BS = read_results(unc_type='BS', ds_name=ds_name, severity=severity, target=target)\n",
    "    results_Conf = read_results(unc_type='Conf', ds_name=ds_name, severity=severity, target=target)\n",
    "    #results_combined = pd.concat([results_Conf, results_DE, results_BS])\n",
    "    results_combined = pd.concat([results_Conf, results_DE])\n",
    "    results_combined['Test set'] = results_combined['Corruption Severity'].replace({\n",
    "        '0': 'CIFAR10',\n",
    "        '1': 'CIFAR10-C (Severity 1)',\n",
    "        '2': 'CIFAR10-C (Severity 2)',\n",
    "        '3': 'CIFAR10-C (Severity 3)',\n",
    "        '4': 'CIFAR10-C (Severity 4)',\n",
    "        '5': 'CIFAR10-C (Highest severity)',\n",
    "    })\n",
    "\n",
    "    plt.rcParams.update({'font.size': 12, 'figure.dpi': 1000})\n",
    "    figure = plt.figure(figsize=(6, 4))\n",
    "    ax = plt.subplot(1, 1, 1)\n",
    "    plot = sns.lineplot(\n",
    "        data=results_combined, x='Validation set quantile', y=target, hue=hue, style='Test set', ax=ax, linewidth=3, \n",
    "        estimator=weighted_mean, errorbar=weighted_sd)\n",
    "    plot.invert_xaxis()\n",
    "    plot.axhline(0.904, c='black', ls='dotted')\n",
    "    sns.despine()\n",
    "    plt.xlim(1.02, 0.6)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig('plots/resnet_{}_quantile_hue{}_sev{}_unc{}_{}.png'.format(ds_name, hue, severity, unc_type, target), bbox_inches='tight')\n",
    "    plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d8069a-5f7c-4f1e-982d-3ed48b8e9c65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure 12a\n",
    "unc_type = 'DE'\n",
    "# Figure 12b\n",
    "#unc_type = 'TTA'\n",
    "sns.set_palette(['#d62728', '#1f77b4'])\n",
    "target = 'Accuracy'\n",
    "style = 'Test set'\n",
    "hue = 'Uncertainty'\n",
    "\n",
    "results_combined = []\n",
    "for severity in range(1, 6):\n",
    "    results_ = read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)\n",
    "    results_['Test set'] = results_['Corruption Severity'].replace({'0': 'CIFAR10', '1': 'CIFAR10-C', '2': 'CIFAR10-C', '3': 'CIFAR10-C', '4': 'CIFAR10-C', '5': 'CIFAR10-C'})\n",
    "    results_combined += [results_]\n",
    "    results_ = read_results(unc_type='Conf', ds_name=ds_name, severity=severity, target=target)\n",
    "    results_['Test set'] = results_['Corruption Severity'].replace({'0': 'CIFAR10', '1': 'CIFAR10-C', '2': 'CIFAR10-C', '3': 'CIFAR10-C', '4': 'CIFAR10-C', '5': 'CIFAR10-C'})\n",
    "    results_combined += [results_]\n",
    "\n",
    "results_combined = pd.concat(results_combined)\n",
    "\n",
    "plt.rcParams.update({'font.size': 12, 'figure.dpi': 1000})\n",
    "figure = plt.figure(figsize=(6, 4))\n",
    "ax = plt.subplot(1, 1, 1)\n",
    "plot = sns.lineplot(\n",
    "    data=results_combined, x='Validation set quantile', y=target, hue=hue, style=style, ax=ax, linewidth=3, \n",
    "    estimator=weighted_mean, errorbar=weighted_sd)\n",
    "plot.invert_xaxis()\n",
    "plot.axhline(0.904, c='black', ls='dotted')\n",
    "sns.despine()\n",
    "plt.xlim(1.02, 0.6)\n",
    "plt.ylim(None, 1.01)\n",
    "plt.tight_layout()\n",
    "plt.savefig('plots/resnet_{}_quantile_hue{}_sev{}_unc{}_{}.png'.format(ds_name, hue, 'ALL', unc_type, target), bbox_inches='tight')\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f8ef7bb-8ac2-439d-bda1-5c0bc2d5b65e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure 11a\n",
    "sns.set_palette(['#d62728', '#1f77b4'])\n",
    "target = 'NLL'\n",
    "unc_type = 'DE'\n",
    "style = 'Test set'\n",
    "hue = 'Uncertainty'\n",
    "\n",
    "results_combined = []\n",
    "for severity in range(1, 6):\n",
    "    results_ = read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)\n",
    "    results_['Test set'] = results_['Corruption Severity'].replace({'0': 'CIFAR10', '1': 'CIFAR10-C', '2': 'CIFAR10-C', '3': 'CIFAR10-C', '4': 'CIFAR10-C', '5': 'CIFAR10-C'})\n",
    "    results_combined += [results_]\n",
    "    results_ = read_results(unc_type='Conf', ds_name=ds_name, severity=severity, target=target)\n",
    "    results_['Test set'] = results_['Corruption Severity'].replace({'0': 'CIFAR10', '1': 'CIFAR10-C', '2': 'CIFAR10-C', '3': 'CIFAR10-C', '4': 'CIFAR10-C', '5': 'CIFAR10-C'})\n",
    "    results_combined += [results_]\n",
    "\n",
    "results_combined = pd.concat(results_combined)\n",
    "\n",
    "plt.rcParams.update({'font.size': 12, 'figure.dpi': 1000})\n",
    "figure = plt.figure(figsize=(6, 4))\n",
    "ax = plt.subplot(1, 1, 1)\n",
    "plot = sns.lineplot(\n",
    "    data=results_combined, x='Validation set quantile', y=target, hue=hue, style=style, ax=ax, linewidth=3, \n",
    "    estimator=weighted_mean, errorbar=weighted_sd)\n",
    "plot.invert_xaxis()\n",
    "sns.despine()\n",
    "plt.xlim(1.02, 0.4)\n",
    "plt.tight_layout()\n",
    "plt.savefig('plots/resnet_{}_quantile_hue{}_sev{}_unc{}_{}.png'.format(ds_name, hue, 'ALL', unc_type, target), bbox_inches='tight')\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "043d3bdd-9368-4a07-927d-629524bc2e5c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}