day_5/5_yingzhen/classification.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ProbAI 2022 BNN Tutorial - classification.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# A Hands-on Tutorial for Bayesian Neural Networks\n",
        "\n",
        "(Part 2: classification)\n",
        "\n",
        "[Yingzhen Li](http://yingzhenli.net/home/en/)\n",
        "\n",
        "(As part of the BNN lecture at [ProbAI 2022](https://probabilistic.ai/))\n",
        "\n",
        "In this tutorial, you will implement various Bayesian neural network methods based on variational inference.\n",
        "\n",
        "We will go through classification tasks to see the applications of uncertainty estimation in practice, including a case study on detecting OOD examples.\n",
        "\n",
        "**How to use this tutorial notebook?**\n",
        "\n",
        "*   Read the descriptions in the text;\n",
        "*   Fill in the missing code whenever you see a block that looks like below:\n",
        "\n",
        "```\n",
        "### beginning of your code ###\n",
        "[insert your own code here]\n",
        "### end of your code ###\n",
        "```\n",
        "There will be hints provided in the code blocks as well to guide you through.\n",
        "\n",
        "**GPU usage**: For this demo, it is recommended to use GPU to accelate training. To do so, click on \"Runtime > Change runtime type\", and select \"GPU\" for the \"Hardware accelerator\" option.\n",
        "\n",
        "Let us set up the required packages, and then, enjoy 😊"
      ],
      "metadata": {
        "id": "A9-xaxhnZGmM"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ZmBUecUY3E4"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "\n",
        "import numpy as np\n",
        "import math\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributions as dist\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "EPS = 1e-5  # define a small constant for numerical stability control"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training BNNs on image classfication tasks\n",
        "\n",
        "In this part you will train a BNN on an image classfication task. More specifically, we will build a Bayesian convolutional neural network (Bayesian CNN) and test a number of approximate posterior inference methods. The dataset in use is MNIST (hand-written digit classification)."
      ],
      "metadata": {
        "id": "P_gLdppUaIlv"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's set up the MNIST dataset first."
      ],
      "metadata": {
        "id": "fodHNzqXe0Kf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# setting up the MNIST dataset\n",
        "transform=transforms.Compose([\n",
        "        transforms.ToTensor(),])\n",
        "train_data = datasets.MNIST('../data', train=True, download=True,\n",
        "                       transform=transform)\n",
        "test_data = datasets.MNIST('../data', train=False,\n",
        "                       transform=transform)\n",
        "train_loader = torch.utils.data.DataLoader(train_data, batch_size=100, shuffle=True)\n",
        "test_loader = torch.utils.data.DataLoader(test_data, batch_size=100, shuffle=False)"
      ],
      "metadata": {
        "id": "g8uOOI_efF8S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Implementing a Bayesian CNN with MC dropout\n",
        "\n",
        "In below we construct a Bayesian CNN with MC dropout and set-up the prediction function with dropout activated."
      ],
      "metadata": {
        "id": "I2ATbBeZay9s"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# construct a Bayesian CNN for MNIST data\n",
        "def make_bayesian_cnn(in_channel, n_class, dropout_prob=0.1, activation='LeakyReLU'):\n",
        "    nonlinearity = getattr(nn, activation)() if isinstance(activation, str) else activation\n",
        "    net = nn.Sequential()\n",
        "    # add in conv layers\n",
        "    channel_sizes = [in_channel, 32, 64, 128]\n",
        "    for i, (c_in, c_out) in enumerate(zip(channel_sizes[:-1], channel_sizes[1:])):\n",
        "        net.add_module(f'Conv{i}', nn.Conv2d(c_in, c_out, kernel_size=3, padding='same'))\n",
        "        if dropout_prob > EPS:  # if setting dropout_prob=0.0 then constructing a deterministic CNN\n",
        "            net.add_module(f'Dropout{i}', nn.Dropout(dropout_prob))\n",
        "        net.add_module(f'Nonlinarity{i}', nonlinearity)\n",
        "        net.add_module(f'Pooling{i}', nn.MaxPool2d(kernel_size=2, ceil_mode=True))\n",
        "    # add in the last fully connected layer\n",
        "    # the image sizes goes from 28*28 -> 14*14 -> 7*7 -> 4*4\n",
        "    net.add_module(f'Flatten', nn.Flatten(start_dim=1, end_dim=-1))\n",
        "    net.add_module(f'Linear', nn.Linear(128*4*4, n_class))\n",
        "    return net\n",
        "\n",
        "# define the prediction function with Monte Carlo sampling using K samples\n",
        "def predict(bnn, x_test, K=1, use_dropout=True, reduce_mean=True):\n",
        "    y_pred = []\n",
        "    if use_dropout:\n",
        "        bnn.train() # this will turn on dropout in prediction time, if applicable\n",
        "    else:\n",
        "        bnn.eval()  # turn off dropout, if applicable\n",
        "    for _ in range(K):\n",
        "        y_pred.append(F.softmax(bnn(x_test), dim=-1))\n",
        "    # shape (K, batch_size, y_dim) or (batch_size, y_dim) if K = 1\n",
        "    y_pred = torch.stack(y_pred, dim=0).squeeze(0)\n",
        "    if reduce_mean and K > 1: # \"ensemble\"\n",
        "        y_pred = y_pred.mean(0)\n",
        "    return y_pred"
      ],
      "metadata": {
        "id": "MgJ1eCUZbT4f"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Adding some helper functions..."
      ],
      "metadata": {
        "id": "wPwN7r3X5W_M"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def to_numpy(x):\n",
        "    return x.detach().cpu().numpy() # convert a torch tensor to a numpy array\n",
        "\n",
        "# helper functions for training\n",
        "def train_step(net, opt, dataloader, device):\n",
        "    logs = []\n",
        "    for _, (x, y) in enumerate(dataloader):\n",
        "        x = x.to(device); y = y.to(device)\n",
        "        opt.zero_grad() # opt is the optimiser\n",
        "        y_logit = net(x)\n",
        "        # training accruacy (on a mini-batch)\n",
        "        pred = y_logit.data.max(1, keepdim=True)[1] # get the index of the max logit\n",
        "        acc = pred.eq(y.data.view_as(pred)).float().cpu().mean()\n",
        "        # training loss\n",
        "        nll = F.nll_loss(F.log_softmax(y_logit, dim=-1), y)\n",
        "        loss = nll  # note that by using AdamW optimizer, weight-decay is implicitly handled there\n",
        "        loss.backward()\n",
        "        opt.step()\n",
        "        logs.append([to_numpy(nll), to_numpy(acc)])\n",
        "    return np.array(logs)\n",
        "\n",
        "# define the training function\n",
        "def train_network(net, opt, dataloader, device, N_epochs=2000, verbose=True):\n",
        "    net.train()\n",
        "    logs = []\n",
        "    for i in range(N_epochs):\n",
        "        logs_epoch = train_step(net, opt, dataloader, device)\n",
        "        if verbose:\n",
        "            print(\"Epoch {}, last mini-batch nll={}, acc={}\".format(i+1, logs_epoch[-1][0], logs_epoch[-1][1]))\n",
        "        logs.append(logs_epoch)\n",
        "    return np.concatenate(logs, axis=0)"
      ],
      "metadata": {
        "id": "qPHRp8RCxtCy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we train a Bayesian CNN using MC-dropout"
      ],
      "metadata": {
        "id": "leKbq9FN5cCn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# constructing and training a Bayesian CNN\n",
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "print(\"using {}\".format(device))\n",
        "dropout_prob = 0.1\n",
        "activation = 'ReLU'\n",
        "dropout_bnn = make_bayesian_cnn(in_channel=1, n_class=10, dropout_prob=dropout_prob, \n",
        "                          activation=activation)\n",
        "dropout_bnn.to(device)\n",
        "print(dropout_bnn)\n",
        "\n",
        "# start training\n",
        "learning_rate = 1e-3\n",
        "weight_decay=0.1\n",
        "N_epochs = 5\n",
        "# note that weight decay (i.e., prior loss) is implicitly handled by AdamW optimizer\n",
        "opt = torch.optim.AdamW(dropout_bnn.parameters(), lr=learning_rate, \n",
        "                        weight_decay=weight_decay)\n",
        "# the training loop starts\n",
        "logs = train_network(dropout_bnn, opt, train_loader, device, N_epochs)\n",
        "\n",
        "# plot the training curve\n",
        "def plot_training_loss(logs, dropout_prob, weight_decay):\n",
        "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))\n",
        "    ax1.plot(np.arange(logs.shape[0]), logs[:, 0], 'r-', label='nll')\n",
        "    ax2.plot(np.arange(logs.shape[0]), logs[:, 1], 'r-', label='acc')\n",
        "    ax1.legend()\n",
        "    ax2.legend()\n",
        "    ax1.set_xlabel('epoch')\n",
        "    ax2.set_xlabel('epoch')\n",
        "    ax1.set_title('p={}, weight_decay={}'.format(dropout_prob, weight_decay))\n",
        "    ax2.set_title('p={}, weight_decay={}'.format(dropout_prob, weight_decay))\n",
        "    plt.show()\n",
        "\n",
        "plot_training_loss(logs, dropout_prob, weight_decay)"
      ],
      "metadata": {
        "id": "97QSPqBgxasT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training is finished, we now evaluate its test accuracy. Remember to turn on dropout in test time and do multiple forward passes."
      ],
      "metadata": {
        "id": "woEZgYdi5iSc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test the BNN\n",
        "def evaluate(model, dataloader, device, K=50):\n",
        "    accuracy = 0\n",
        "    for x, y in test_loader:\n",
        "        x = x.to(device); y = y.to(device)\n",
        "        # the \"predict\" function: by default dropout is on and we use K > 1 MC samples\n",
        "        y_pred = predict(model, x, K, reduce_mean=True)\n",
        "        pred = y_pred.data.max(1, keepdim=True)[1] # get the index of the max probability\n",
        "        accuracy += pred.eq(y.data.view_as(pred)).float().cpu().sum()\n",
        "    accuracy = accuracy / len(dataloader.dataset) * 100 # accuracy in percentage\n",
        "    return to_numpy(accuracy)\n",
        "\n",
        "accuracy = evaluate(dropout_bnn, test_loader, device, K=5)\n",
        "print('Test Accuracy: {}%'.format(accuracy))"
      ],
      "metadata": {
        "id": "x5Tac3ju1h5U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## MFVI for Bayesian CNN\n",
        "\n",
        "We will also train a Bayesian CNN with MFVI approximations. For this we will use the ```Bayesianize``` package. ```Bayesianize``` is a lightweight pytorch package which can convert any MLPs or CNNs to their \"Bayesian\" version using variational inference.\n",
        "\n",
        "First we clone and import the package."
      ],
      "metadata": {
        "id": "-pq1rMfHCwkO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/microsoft/bayesianize.git"
      ],
      "metadata": {
        "id": "0ez5JMLVDTN4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.path.append('bayesianize/')\n",
        "import bnn"
      ],
      "metadata": {
        "id": "vbEzUhvoDXvT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We prepare some helper training functions in accordance to the usage of the ```Bayesianize``` package. The idea is similar to the ones in the regression tutorial: we need to compute both the \"data losses\" (i.e., the negative log-likelihood terms) as well as the KL terms."
      ],
      "metadata": {
        "id": "BZJiHXvB60Fh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# helper functions for training\n",
        "def train_bnn_step(net, opt, dataloader, device, beta):\n",
        "    logs = []\n",
        "    N_data = len(dataloader.dataset)\n",
        "    for _, (x, y) in enumerate(dataloader):\n",
        "        x = x.to(device); y = y.to(device)\n",
        "        opt.zero_grad() # opt is the optimiser\n",
        "        y_logit = net(x)\n",
        "        # training accruacy (on a mini-batch)\n",
        "        pred = y_logit.data.max(1, keepdim=True)[1] # get the index of the max logit\n",
        "        acc = pred.eq(y.data.view_as(pred)).float().cpu().mean()\n",
        "        # training loss\n",
        "        nll = F.nll_loss(F.log_softmax(y_logit, dim=-1), y)\n",
        "        kl = sum(m.kl_divergence() for m in net.modules() if hasattr(m, \"kl_divergence\"))\n",
        "        loss = N_data * nll + beta * kl\n",
        "        loss.backward()\n",
        "        opt.step()\n",
        "        logs.append([to_numpy(nll), to_numpy(acc), to_numpy(kl)])\n",
        "    return np.array(logs)\n",
        "\n",
        "# define the training function\n",
        "def train_bnn(net, opt, dataloader, device, N_epochs=2000, beta=1.0, verbose=True):\n",
        "    net.train()\n",
        "    logs = []\n",
        "    for i in range(N_epochs):\n",
        "        logs_epoch = train_bnn_step(net, opt, dataloader, device, beta)\n",
        "        if verbose:\n",
        "            print(\"Epoch {}, last mini-batch nll={}, acc={}, kl={}\".format(\n",
        "                i+1, logs_epoch[-1][0], logs_epoch[-1][1], logs_epoch[-1][2]))\n",
        "        logs.append(logs_epoch)\n",
        "    return np.concatenate(logs, axis=0)"
      ],
      "metadata": {
        "id": "hsprjWfBF4Sd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we train a Bayesian CNN with MFVI as for approximate posterior inference."
      ],
      "metadata": {
        "id": "IQbuHwZZ7PjY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "activation = 'ReLU'\n",
        "# use dropout_prob = 0.0 effectively constructs a deterministic CNN\n",
        "mfvi_bnn = make_bayesian_cnn(in_channel=1, n_class=10, dropout_prob=0.0, \n",
        "                             activation=activation)\n",
        "# use the \"Bayesianize\" package to make a Bayesian CNN with MFVI\n",
        "# inference='ffg' means using fully-factorised Gaussian for VI, i.e., MFVI\n",
        "# see https://github.com/microsoft/bayesianize/blob/main/bnn/nn/mixins/variational/ffg.py\n",
        "bnn.bayesianize_(mfvi_bnn, inference='ffg', init_sd=0.02)\n",
        "mfvi_bnn.to(device)\n",
        "print(mfvi_bnn)\n",
        "\n",
        "# start training\n",
        "learning_rate = 1e-3\n",
        "N_epochs = 5\n",
        "beta = 1.0\n",
        "opt = torch.optim.Adam(mfvi_bnn.parameters(), lr=learning_rate)\n",
        "# the training loop starts\n",
        "logs = train_bnn(mfvi_bnn, opt, train_loader, device, N_epochs, beta)\n",
        "\n",
        "# plot the training curve\n",
        "def plot_training_loss(logs, title):\n",
        "    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 4))\n",
        "    ax1.plot(np.arange(logs.shape[0]), logs[:, 0], 'r-', label='nll')\n",
        "    ax2.plot(np.arange(logs.shape[0]), logs[:, 1], 'r-', label='acc')\n",
        "    ax3.plot(np.arange(logs.shape[0]), logs[:, 2], 'r-', label='KL')\n",
        "    for ax in [ax1, ax2, ax3]:\n",
        "        ax.legend()\n",
        "        ax.set_xlabel('epoch')\n",
        "        ax1.set_title(title)\n",
        "    plt.show()\n",
        "\n",
        "plot_training_loss(logs, title='MFVI')"
      ],
      "metadata": {
        "id": "H5WrGFI6D53-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training is finished, then we evaluate the accuracy of MFVI Bayesian CNN."
      ],
      "metadata": {
        "id": "otYKHXCKKtqm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "accuracy = evaluate(mfvi_bnn, test_loader, device, K=5)\n",
        "print('Test Accuracy: {}%'.format(accuracy))"
      ],
      "metadata": {
        "id": "sEjZezhBIXGt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Case study 2: Detecting adversarial attacks\n",
        "\n",
        "Neural networks has been shown to be vulnerable to adversarial attacks. In this case study, we will see whether BNNs can be more robust or not, as well as how to use uncertainty measures from BNNs to detect adversarial examples.\n",
        "\n",
        "\n",
        "But first let's set-up the helper functions for the adversarial attack. Here we consider the $L_\\infty$-norm projected gradient descent (PGD) attack [(Madry et al. ICLR 2018)](https://arxiv.org/abs/1706.06083) which is one of the strongest attack considerred in adversarial robustness literature."
      ],
      "metadata": {
        "id": "jzYt8P6t-i5R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# PGD attack for one step\n",
        "def pgd_one_step(x, x_grad, x_clean, epsilon, lr, x_min, x_max):\n",
        "    # use scaled element-wise sign of the data gradient for attack\n",
        "    x_adv = x + lr * x_grad.sign()\n",
        "    x_adv = torch.clamp(x_adv, min=x_min, max=x_max)  # make sure x is valid\n",
        "    # project the attack back to the epsilon-ball (L-inf norm)\n",
        "    x_adv = x_clean + torch.clamp(x_adv - x_clean, min=-epsilon, max=epsilon)\n",
        "    return x_adv\n",
        "\n",
        "def pgd_attack(x, y, model, predict_func, epsilon, T, lr, x_min, x_max):\n",
        "    x_adv = x.detach().clone()\n",
        "    x_adv.requires_grad = True  # we need to compute gradients w.r.t. the input\n",
        "    # perform PGD attack for T steps\n",
        "    for t in range(T):\n",
        "        # Forward pass the data through the model and get loss\n",
        "        y_pred = predict_func(x_adv, reduce_mean=True) # return prediction probabilities\n",
        "        loss = F.nll_loss(torch.log(torch.clamp(y_pred, min=EPS)), y)\n",
        "        # One-step PGD using gradients of loss w.r.t. input in backward pass\n",
        "        if type(model) is list:\n",
        "            for net in model:\n",
        "                net.zero_grad()\n",
        "        else:\n",
        "            model.zero_grad()\n",
        "        loss.backward()\n",
        "        x_adv.data = pgd_one_step(x_adv.data, x_adv.grad.data, x, epsilon, lr, x_min, x_max)\n",
        "    return x_adv"
      ],
      "metadata": {
        "id": "cNuL8oCJ_dC3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We first test the PGD attack on a deterministic network (as to validate the implementation). For this deterministic network, we just need to turn off dropout during test time."
      ],
      "metadata": {
        "id": "YfD0zwEO9ZW2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# testing adversarial attacks\n",
        "# we will use 100 test samples as examples\n",
        "x_test, y_test = next(iter(test_loader))\n",
        "x_test = torch.tensor(x_test, device=device)\n",
        "y_test = torch.tensor(y_test, device=device)\n",
        "\n",
        "# define some helper functions for the attack\n",
        "def run_attack(model, x_test, y_test, predict_func, epsilon_list, T):\n",
        "    acc_adv = []\n",
        "    pred_adv = []\n",
        "    prob_adv = []\n",
        "    x_adv = []\n",
        "    for epsilon in epsilon_list:\n",
        "        lr = epsilon / T * 1.5\n",
        "        x_adv_ = pgd_attack(x_test, y_test, model, predict_func, \n",
        "                           epsilon, T, lr, x_min=0.0, x_max=1.0)\n",
        "        # get the accuracy on crafted adversarial examples\n",
        "        y_pred = predict_func(x_adv_, reduce_mean=True)\n",
        "        pred = y_pred.data.max(1, keepdim=True)[1] # get the index of the max probability\n",
        "        acc_adv.append(to_numpy(pred.eq(y_test.data.view_as(pred)).float().mean()))\n",
        "        pred_adv.append(to_numpy(pred))\n",
        "        prob_adv.append(to_numpy(y_pred))\n",
        "        x_adv.append(to_numpy(x_adv_))\n",
        "    return np.array(acc_adv), np.array(pred_adv), np.array(prob_adv), np.array(x_adv)\n",
        "\n",
        "# helper function for visualization\n",
        "def visualize(images, epsilons, pred_clean, pred_adv):\n",
        "    cnt = 0\n",
        "    plt.figure(figsize=(8, 2 * len(epsilons)))\n",
        "    for i in range(len(epsilons)):\n",
        "        for j in range(len(images[i])):\n",
        "            cnt += 1\n",
        "            plt.subplot(len(epsilons),len(images[0]),cnt)\n",
        "            plt.xticks([], [])\n",
        "            plt.yticks([], [])\n",
        "            if j == 0:\n",
        "                plt.ylabel(\"Eps: {:.2f}\".format(epsilons[i]), fontsize=14)\n",
        "            plt.title(\"{} -> {}\".format(pred_clean[j], pred_adv[i, j]))\n",
        "            plt.imshow(images[i, j][0], cmap=\"gray\")\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "# we first test a deterministic NN's robustness, by turning off dropout\n",
        "predict_deterministic = lambda x, reduce_mean: predict(dropout_bnn, x, K=1, use_dropout=False, \n",
        "                                                       reduce_mean=reduce_mean)\n",
        "\n",
        "def test_robustness(model, predict_func, x_test, y_test):\n",
        "    # first test clean accuracy\n",
        "    y_pred = predict_func(x_test, reduce_mean=True)\n",
        "    entropy = -(y_pred * torch.log(y_pred)).sum(-1).mean()\n",
        "    print(\"entropy (clean)\", to_numpy(entropy))\n",
        "    pred_clean = y_pred.data.max(1, keepdim=True)[1] # get the index of the max probability\n",
        "    acc_clean = to_numpy(pred_clean.eq(y_test.data.view_as(pred_clean)).float().mean())\n",
        "    pred_clean = to_numpy(pred_clean)\n",
        "    print(\"clean accruacy for this batch: {}%\".format(acc_clean * 100))\n",
        "    \n",
        "    # now run the attack\n",
        "    T = 10\n",
        "    epsilon_list = np.arange(0.05, 0.41, 0.05)\n",
        "    acc_adv, pred_adv, prob_adv, x_adv = run_attack(model, x_test, y_test, predict_func, epsilon_list, T)\n",
        "    print(\"adv accruacy for this batch (%): {}\".format(acc_adv * 100))\n",
        "    # visualise the adversarial examples\n",
        "    visualize(x_adv[:, :5], epsilon_list, pred_clean[:5], pred_adv[:, :5])\n",
        "    return x_adv\n",
        "\n",
        "x_adv = test_robustness(dropout_bnn, predict_deterministic, x_test, y_test)"
      ],
      "metadata": {
        "id": "2i8dSWHmGcDc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now testing the MC-dropout Bayesian CNN's robustness. \n",
        "\n",
        "In this case we will also look at the uncertainty produced by the BNN on the adversarial examples. Specifically, you will implement:\n",
        "\n",
        "*   total uncertainty: as the total entropy\n",
        "*   aleatoric uncertainty: as the conditional entropy\n",
        "\n",
        "Epistemic uncertainty is the mutual information which is essentially total entropy minus conditional entropy."
      ],
      "metadata": {
        "id": "5XV-VU6IQqIi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# now we test Bayesian CNN's robustness, by using MC-dropout\n",
        "predict_mcdropout = lambda x, reduce_mean: predict(dropout_bnn, x, K=5, use_dropout=True,\n",
        "                                                   reduce_mean=reduce_mean)\n",
        "x_adv = test_robustness(dropout_bnn, predict_mcdropout, x_test, y_test)\n",
        "\n",
        "# also compute uncertainty measures\n",
        "### beginning of your code ###\n",
        "# Hints: for below two functions, the input y_prob_samples contains classification \n",
        "# probability vector from the BNN using weights sampled from the approximate posterior. \n",
        "# It should have shape y_prob_samples.shape = (K, batch_size, n_class),\n",
        "# with K the number of MC samples and n_class = 10 in MNIST case.\n",
        "# We would like to compute the entropy/conditional entropy using y_prob_samples,\n",
        "# and the returned tensor should have shape (batch_size,).\n",
        "\n",
        "def total_entropy(y_prob_samples):\n",
        "    entropy =\n",
        "    return entropy\n",
        "\n",
        "def conditional_entropy(y_prob_samples):\n",
        "    cond_entropy = \n",
        "    return cond_entropy\n",
        "\n",
        "### end of your code ###\n",
        "\n",
        "def compute_uncertainty(predict_func, x_adv, reduce_mean=True):\n",
        "    logs = []\n",
        "    for i in range(x_adv.shape[0]):\n",
        "        x = torch.tensor(x_adv[i], device=device)\n",
        "        y_prob_samples = predict_func(x, reduce_mean=False)\n",
        "        entropy = total_entropy(y_prob_samples)\n",
        "        cond_entropy = conditional_entropy(y_prob_samples)\n",
        "        if reduce_mean:\n",
        "            entropy = entropy.mean(); cond_entropy = cond_entropy.mean()\n",
        "        mutual_info = entropy - cond_entropy\n",
        "        logs.append([to_numpy(entropy), to_numpy(cond_entropy), to_numpy(mutual_info)])\n",
        "    return np.array(logs)\n",
        "\n",
        "logs = compute_uncertainty(predict_mcdropout, x_adv)\n",
        "print(\"total uncertainty (entropy):\", logs[:, 0])\n",
        "print(\"aleatoric uncertainty (cond. entropy):\", logs[:, 1])\n",
        "print(\"epistemic uncertainty (mutual info):\", logs[:, 2])"
      ],
      "metadata": {
        "id": "yrCJ71w3RA-_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And similarly for the MFVI Bayesian CNN, test its adversarial robustness."
      ],
      "metadata": {
        "id": "G6910FAQ9gOf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# now we test Bayesian CNN's robustness, by using MFVI\n",
        "predict_mfvi = lambda x, reduce_mean: predict(mfvi_bnn, x, K=5, reduce_mean=reduce_mean)\n",
        "x_adv = test_robustness(mfvi_bnn, predict_mfvi, x_test, y_test)\n",
        "# also compute epistemic uncertainty\n",
        "logs = compute_uncertainty(predict_mfvi, x_adv)\n",
        "print(\"total uncertainty (entropy):\", logs[:, 0])\n",
        "print(\"aleatoric uncertainty (cond. entropy):\", logs[:, 1])\n",
        "print(\"epistemic uncertainty (mutual info):\", logs[:, 2])"
      ],
      "metadata": {
        "id": "KBVKHfvdCRkR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Can ensembles help?\n",
        "\n",
        "From the above examples we see that the trained BNNs help a bit in adversarial robustness but not too much, especially for the MC-dropout one. From the printed uncertainty measures, we suspect the trained networks may be still not uncertain enough on the adversarial examples.\n",
        "\n",
        "Next we will train an Ensemble of MFVI-BNNs. This means we will train multiple MFVI-BNNs independently with different random initialisation of the variational parameters, and then ensemble them together to get the final model. Although diversity is not explicitly encouraged, the hope is that the randomness in initialisation and SGD steps would allow different independent runs to result in different local optima, resulting in a more diversed set of neural networks.\n",
        "\n",
        "Please run the code in below block, it might take a few minutes to train the MFVI-BNNs in the ensemble."
      ],
      "metadata": {
        "id": "-XNhVvJnf6NG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# train an ensemble of MFVI networks (independently) and see what happens\n",
        "mfvi_bnn_list = [mfvi_bnn]\n",
        "N_ensemble = 5\n",
        "activation = 'ReLU'\n",
        "learning_rate = 1e-3\n",
        "N_epochs = 5\n",
        "beta = 1.0\n",
        "\n",
        "for i in range(1, N_ensemble):\n",
        "    print(\"train a new MFVI CNN (network {})\".format(i+1))\n",
        "    net = make_bayesian_cnn(in_channel=1, n_class=10, dropout_prob=0.0, \n",
        "                            activation=activation)\n",
        "    bnn.bayesianize_(net, inference='ffg', init_sd=0.02)\n",
        "    net.to(device)\n",
        "    print(net)\n",
        "    # start training\n",
        "    opt = torch.optim.Adam(net.parameters(), lr=learning_rate)\n",
        "    logs = train_bnn(net, opt, train_loader, device, N_epochs, beta)\n",
        "    mfvi_bnn_list.append(net)"
      ],
      "metadata": {
        "id": "FdvAe0Smf8w6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training is now finished, let's see how the ensemble performs on clean data."
      ],
      "metadata": {
        "id": "Uct-1Us2MuzT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test the ensemble on clean test data\n",
        "def predict_ensemble(x_test, K=5, reduce_mean=True):\n",
        "    y_prob_samples = []\n",
        "    for net in mfvi_bnn_list:\n",
        "        y_prob_samples.append(predict(net, x_test, K=K, reduce_mean=False))\n",
        "    y_prob_samples = torch.concat(y_prob_samples, axis=0)\n",
        "    if reduce_mean:\n",
        "        return y_prob_samples.mean(0)\n",
        "    else:\n",
        "        return y_prob_samples\n",
        "\n",
        "accuracy = 0\n",
        "K = 5\n",
        "for x, y in test_loader:\n",
        "    x = x.to(device); y = y.to(device)\n",
        "    y_pred = predict_ensemble(x, K)\n",
        "    pred = y_pred.data.max(1, keepdim=True)[1] # get the index of the max probability\n",
        "    accuracy += pred.eq(y.data.view_as(pred)).float().cpu().sum()\n",
        "accuracy = to_numpy(accuracy / len(test_loader.dataset) * 100) # accuracy in percentage\n",
        "print('Test Accuracy: {}%'.format(accuracy))"
      ],
      "metadata": {
        "id": "pg3t3s-Gjsu-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we perform PGD attacks on the ensemble BNN."
      ],
      "metadata": {
        "id": "4vuukXBVMzfb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test for adversarial attacks\n",
        "x_adv = test_robustness(mfvi_bnn_list, predict_ensemble, x_test, y_test)\n",
        "# also compute epistemic uncertainty\n",
        "logs = compute_uncertainty(predict_ensemble, x_adv)\n",
        "print(\"total uncertainty (entropy):\", logs[:, 0])\n",
        "print(\"aleatoric uncertainty (cond. entropy):\", logs[:, 1])\n",
        "print(\"epistemic uncertainty (mutual info):\", logs[:, 2])"
      ],
      "metadata": {
        "id": "m4_eVAjYjpxJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Quantitative metric for detection results\n",
        "\n",
        "The ensemble BNN above makes a bit of improvement on robustness but not too much, however, in terms of uncertainty measures, the ensemble becomes more uncertain on the adversarial inputs. But at the same time, the total entropy on clean inputs remains relatively close to the single MFVI-BNN case. This is helpful as it will be easier to separate clean/adversarial examples by looking at the uncertainty measures and picking appropriate thresholds for dectection.\n",
        "\n",
        "In below blocks, we will quantify the quality of uncertainty by considering its effectiveness as a tool for adversarial example detection. The idea is, for a given uncertainty measure (e.g., total entropy):\n",
        "\n",
        "*   for a given threshold, predict the input as \"adversarial\" if uncertainty > threshold;\n",
        "*   search for the best threshold so that it returns the best true positive rate (TPR) of detecting adversarial example, while maintaining the false positive rate (FPR) to be below 5%.\n",
        "\n"
      ],
      "metadata": {
        "id": "cAMnjnwM9dkH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# below are helper functions for the adversarial example detection test\n",
        "\n",
        "def get_uncertainty_adv_examples(model, predict_func, x_test, y_test, epsilon_list, T=10):\n",
        "    # first get uncertainties on clean examples, return tensor of shape (3, batch_size)\n",
        "    uncertainty_clean = compute_uncertainty(predict_func, x_test.unsqueeze(0), reduce_mean=False)[0]\n",
        "    # now run the attack\n",
        "    acc_adv, pred_adv, prob_adv, x_adv = run_attack(model, x_test, y_test, predict_func, epsilon_list, T)\n",
        "    print(\"adv accruacy for this batch (%): {}\".format(acc_adv * 100))\n",
        "    # get uncertainties on adversarial examples, return tensor of shape (len(epsilon_list, 3, batch_size))\n",
        "    uncertainty_adv = compute_uncertainty(predict_func, x_adv, reduce_mean=False)\n",
        "    return uncertainty_clean, uncertainty_adv\n",
        "\n",
        "def compute_detection_error(thresholds, uncertainty, label):\n",
        "    # compute the detection error\n",
        "    # here if uncertainty > threshold then we say it is an OOD example\n",
        "    # where for label, 0 means in-distribution, 1 means OOD\n",
        "    # we assume the first half of the labels are 0 and the second half of the labels are 1\n",
        "    tpr = []; fpr = []\n",
        "    N = int(label.shape[0] / 2)\n",
        "    for t in thresholds:\n",
        "        pred = np.maximum(np.sign(uncertainty - t), 0)\n",
        "        fpr.append(1 - np.mean(np.where(pred[:N] == label[:N], 1, 0)))\n",
        "        tpr.append(np.mean(np.where(pred[N:] == label[N:], 1, 0)))\n",
        "    return fpr, tpr\n",
        "\n",
        "def run_detect(model, x_test, y_test, predict_func, epsilon_list, min_fpr = 0.05):\n",
        "    thresholds = np.arange(0.0, np.log(10), 0.05)\n",
        "    fpr_total = 0.0; tpr_total = 0.0\n",
        "    N_data = x_test.shape[0]; batch_size = 100; N_batches = int(N_data / batch_size)\n",
        "    for n in range(N_batches):\n",
        "        x = x_test[n*batch_size:(n+1)*batch_size].to(device)\n",
        "        y = y_test[n*batch_size:(n+1)*batch_size].to(device)\n",
        "        uncertainty_clean, uncertainty_adv = get_uncertainty_adv_examples(model, predict_func, x, y, epsilon_list)\n",
        "        # we assume the first half of the labels are 0 and the second half of the labels are 1\n",
        "        label = np.concatenate([np.zeros(x.shape[0]), np.ones(x.shape[0])])\n",
        "        # len(uncertainty_adv) == len(epsilon_list)\n",
        "        fpr = np.zeros((len(uncertainty_adv), 3, len(thresholds)))\n",
        "        tpr = np.zeros((len(uncertainty_adv), 3, len(thresholds)))\n",
        "        for i in range(len(uncertainty_adv)):\n",
        "            for j in range(3):\n",
        "                uncertainty = np.concatenate([uncertainty_clean[j], uncertainty_adv[i][j]])\n",
        "                fpr[i, j], tpr[i, j] = compute_detection_error(thresholds, uncertainty, label)\n",
        "        fpr_total = fpr_total + fpr * batch_size / N_data\n",
        "        tpr_total = tpr_total + tpr * batch_size / N_data\n",
        "    # now we compute the best tpr given that the corresponding fpr <= min_fpr\n",
        "    mask = np.where(fpr_total <= min_fpr, 1, 0)\n",
        "    tpr_best = np.max(tpr_total * mask, axis=-1)\n",
        "    # plot the result\n",
        "    plt.subplots(1, 1, figsize=(5, 4))\n",
        "    plt.plot(epsilon_list, tpr_best[:, 0], 'k-', label='total entropy')\n",
        "    plt.plot(epsilon_list, tpr_best[:, 1], 'b-', label='cond entropy')\n",
        "    plt.plot(epsilon_list, tpr_best[:, 2], 'r-', label='mutual info')\n",
        "    plt.legend()\n",
        "    plt.xlabel('epsilon')\n",
        "    plt.title('Best TPR (with FPR <= %.2f)' % min_fpr)\n",
        "    plt.show()\n",
        "    return tpr_best"
      ],
      "metadata": {
        "id": "qXcybj_-Qpu1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can test for adversarial detection. To do so, we randomly sample 1000 (10*100) datapoints from the test data."
      ],
      "metadata": {
        "id": "dC0_tU4fmdJ9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# get N_batch*100 test images to test detection\n",
        "count_batch = 0\n",
        "N_batch = 10\n",
        "x_clean = []; y_clean = []\n",
        "for x, y in test_loader:\n",
        "    x_clean.append(x); y_clean.append(y)\n",
        "    count_batch += 1\n",
        "    if count_batch >= N_batch:\n",
        "        break\n",
        "x_clean = torch.concat(x_clean, dim=0); y_clean = torch.concat(y_clean, dim=0)"
      ],
      "metadata": {
        "id": "lB_VDIV6OuiK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "First we evaluate the MC-dropout method."
      ],
      "metadata": {
        "id": "iPx7CvK-OvdV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test detection for the MC-dropout net\n",
        "epsilon_list = np.arange(0.05, 0.41, 0.05)\n",
        "tpr_dropout = run_detect(dropout_bnn, x_clean, y_clean, predict_mcdropout, epsilon_list)"
      ],
      "metadata": {
        "id": "u_7dewHElx12"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we evaluate the MFVI Bayesian CNN."
      ],
      "metadata": {
        "id": "JWjFNv08mjp3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test detection for the mfvi net\n",
        "tpr_mfvi = run_detect(mfvi_bnn, x_clean, y_clean, predict_mfvi, epsilon_list)"
      ],
      "metadata": {
        "id": "bSNKODDHl09g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Finally we evaluate the ensemble MFVI Bayesian CNN."
      ],
      "metadata": {
        "id": "ZXgUsw6PmmkP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# test detection for the ensemble method\n",
        "tpr_ensemble = run_detect(mfvi_bnn_list, x_clean, y_clean, predict_ensemble, epsilon_list)"
      ],
      "metadata": {
        "id": "OnQzSgaamvML"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Hopefully you can see that the ensemble BNN performs significantly better than the other two BNNs (MC-dropout and single MFVI-BNN) in terms of detecting adversarial examples.\n",
        "\n",
        "There are ways to improve the ensemble BNN further in terms of both robustness and detection to adversarial examples, but they are beyond the scope of this tutorial."
      ],
      "metadata": {
        "id": "lVrAAQcGPLo1"
      }
    }
  ]
}