ensemble_loss_landscape/cifar10_medium_cnn_experiments.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WU--cjOphQKk"
      },
      "outputs": [],
      "source": [
        "  # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.\n",
        "\n",
        "  #  Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "  #  you may not use this file except in compliance with the License.\n",
        "  #  You may obtain a copy of the License at\n",
        "\n",
        "  #      http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "  #  Unless required by applicable law or agreed to in writing, software\n",
        "  #  distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "  #  See the License for the specific language governing permissions and\n",
        "  #  limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MGjUZ2Q55OWM"
      },
      "source": [
        "#Ensemble and Subspace Sampling Experiments on CIFAR10\n",
        "\n",
        "   \u003c!-- Copyright 2020 Stanislav Fořt, Huiyi Hu, Balaji Lakshminarayanan\n",
        "\n",
        "   Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "   you may not use this file except in compliance with the License.\n",
        "   You may obtain a copy of the License at\n",
        "\n",
        "       http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "   Unless required by applicable law or agreed to in writing, software\n",
        "   distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "   See the License for the specific language governing permissions and\n",
        "   limitations under the License. --\u003e\n",
        "\n",
        "\n",
        "\n",
        "This notebook illustrates the CIFAR-10 experiments in the paper: \n",
        "\n",
        "[Deep Ensembles: A Loss Landscape Perspective](https://arxiv.org/abs/1912.02757) by Stanislav Fort, Huiyi Hu and Balaji Lakshminarayanan\n",
        "\n",
        "\n",
        "These experiments investigate the effects of ensembling and variational Bayesian methods, please see the paper for more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U5sMeD_4pTe2"
      },
      "source": [
        "# Setting up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "code",
        "colab": {},
        "colab_type": "code",
        "id": "bUXPnWQIEKJh"
      },
      "outputs": [],
      "source": [
        "%tensorflow_version 1.x\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib as mpl\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "from tensorflow.contrib import layers\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "from scipy.special import softmax\n",
        "from matplotlib import patches as mpatch\n",
        "\n",
        "# Plot Style\n",
        "mpl.style.use('seaborn-colorblind')\n",
        "mpl.rcParams.update({\n",
        "    'font.size': 14,\n",
        "    'lines.linewidth': 2,\n",
        "    'figure.figsize': (6, 6 / 1.61)\n",
        "})\n",
        "mpl.rcParams['grid.color'] = 'k'\n",
        "mpl.rcParams['grid.linestyle'] = ':'\n",
        "mpl.rcParams['grid.linewidth'] = 0.5\n",
        "mpl.rcParams['lines.markersize'] = 6\n",
        "mpl.rcParams['lines.marker'] = None\n",
        "mpl.rcParams['axes.grid'] = True\n",
        "\n",
        "DEFAULT_FONTSIZE = 13\n",
        "mpl.rcParams.update({\n",
        "    'font.size': DEFAULT_FONTSIZE,\n",
        "    'lines.linewidth': 2,\n",
        "    'legend.fontsize': DEFAULT_FONTSIZE,\n",
        "    'axes.labelsize': DEFAULT_FONTSIZE,\n",
        "    'xtick.labelsize': DEFAULT_FONTSIZE,\n",
        "    'ytick.labelsize': DEFAULT_FONTSIZE,\n",
        "    'figure.figsize': (7, 7.0 / 1.4)\n",
        "})\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8_a0OnCQrkFu"
      },
      "source": [
        "#Getting the datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AlG96jyJr8cW"
      },
      "outputs": [],
      "source": [
        "def give_me_data():\n",
        "  print(\"Reading CIFAR10\")\n",
        "  cifar_data = {}\n",
        "  N_val = 500\n",
        "\n",
        "  # Construct a tf.data.Dataset\n",
        "  ds_train, ds_test = tfds.load(\n",
        "      name=\"cifar10\", split=[\"train\", \"test\"], batch_size=-1)\n",
        "\n",
        "  numpy_train = tfds.as_numpy(ds_train)\n",
        "  x_train_raw, y_train_raw = numpy_train[\"image\"], numpy_train[\"label\"]\n",
        "\n",
        "  numpy_test = tfds.as_numpy(ds_test)\n",
        "  x_test_raw, y_test_raw = numpy_test[\"image\"], numpy_test[\"label\"]\n",
        "\n",
        "  N_train = x_train_raw.shape[0] - N_val\n",
        "\n",
        "  X_train = x_train_raw[:N_train]\n",
        "  y_train = y_train_raw[:N_train]\n",
        "  X_val = x_train_raw[N_train:N_train + N_val]\n",
        "  y_val = y_train_raw[N_train:N_train + N_val]\n",
        "  X_test = x_test_raw\n",
        "  y_test = y_test_raw\n",
        "\n",
        "  Hn = 32\n",
        "  Wn = 32\n",
        "  Cn = 3\n",
        "  cifar_data['Hn'] = Hn\n",
        "  cifar_data['Wn'] = Wn\n",
        "  cifar_data['Cn'] = Cn\n",
        "  cifar_data['classes'] = 10\n",
        "\n",
        "  cifar_data['X_train'] = X_train.reshape([-1, Hn, Wn, Cn])\n",
        "  cifar_data['X_val'] = X_val.reshape([-1, Hn, Wn, Cn])\n",
        "  cifar_data['X_test'] = X_test.reshape([-1, Hn, Wn, Cn])\n",
        "\n",
        "  cifar_data['y_train'] = y_train.reshape([-1])\n",
        "  cifar_data['y_val'] = y_val.reshape([-1])\n",
        "  cifar_data['y_test'] = y_test.reshape([-1])\n",
        "  return cifar_data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "height": 35
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 16490,
          "status": "ok",
          "timestamp": 1596259193156,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "P9ius35wN76P",
        "outputId": "2f9f551f-b17c-4e0d-8442-b4749f18f33c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Reading CIFAR10\n"
          ]
        }
      ],
      "source": [
        "cifar_ds = give_me_data()\n",
        "X_train = cifar_ds['X_train']\n",
        "y_train = cifar_ds['y_train']\n",
        "X_val = cifar_ds['X_val']\n",
        "y_val = cifar_ds['y_val']\n",
        "X_test = cifar_ds['X_test']\n",
        "y_test = cifar_ds['y_test']\n",
        "Hn = cifar_ds['Hn']\n",
        "Wn = cifar_ds['Wn']\n",
        "Cn = cifar_ds['Cn']\n",
        "classes = cifar_ds['classes']\n",
        "\n",
        "N_train = X_train.shape[0]\n",
        "N_val = X_val.shape[0]\n",
        "N_test = X_test.shape[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DsEjVGc_pqAL"
      },
      "source": [
        "# Defining a CNN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "fWU-qKjRvvKV"
      },
      "outputs": [],
      "source": [
        "# Builds a CNN network and returns various graph nodes for training and\n",
        "# evaluation.\n",
        "# In the 'eval' part of the graph, placeholders for network weights (Ws \u0026 bs)\n",
        "# are created to facilitate inferencing at given weights (e.g. from subspace\n",
        "# sampling).\n",
        "def multilayer_CNN(X_ph_in,\n",
        "                   y_ph_in,\n",
        "                   dropout_rate_ph_in=None,\n",
        "                   filter_sizes=(3, 3, 3, 3),\n",
        "                   pools=(2, 2, 2, 2),\n",
        "                   channels=(32, 64, 128, 256),\n",
        "                   classes=10):\n",
        "  net_hooks = {\n",
        "      'weights': {},\n",
        "      'placeholder': {},\n",
        "      'train_hook': {},\n",
        "      'train_output': {},\n",
        "      'eval_output': {}\n",
        "  }\n",
        "  f_nonlin = tf.nn.relu\n",
        "  with tf.variable_scope('train'):\n",
        "    a = X_ph_in\n",
        "    Ws = []\n",
        "    bs = []\n",
        "    for i in range(len(filter_sizes)):\n",
        "      _, _, _, Cnow = a.get_shape().as_list()\n",
        "      W = tf.get_variable(\n",
        "          'Wconv' + str(i),\n",
        "          shape=[filter_sizes[i], filter_sizes[i], Cnow, channels[i]],\n",
        "          initializer=layers.xavier_initializer(),\n",
        "          trainable=True)\n",
        "      b = tf.get_variable(\n",
        "          'bconv' + str(i),\n",
        "          shape=[1, 1, channels[i]],\n",
        "          initializer=layers.xavier_initializer(),\n",
        "          trainable=True)\n",
        "      Ws.append(W)\n",
        "      bs.append(b)\n",
        "\n",
        "      h = tf.nn.conv2d(a, W, strides=[1, 1, 1, 1], padding='SAME') + b\n",
        "      h = tf.nn.dropout(h, rate=dropout_rate_ph_in)\n",
        "      h = tf.nn.max_pool(\n",
        "          h,\n",
        "          ksize=[1, pools[i], pools[i], 1],\n",
        "          strides=[1, pools[i], pools[i], 1],\n",
        "          padding='SAME')\n",
        "      if i \u003c len(filter_sizes) - 1:\n",
        "        a = f_nonlin(h)\n",
        "      else:\n",
        "        a = h\n",
        "      _, H_out, W_out, C_out = a.get_shape().as_list()\n",
        "    _, height_final, width_final, channels_final = a.get_shape().as_list()\n",
        "\n",
        "    # Final fully connected layer.\n",
        "    W_final = tf.get_variable(\n",
        "        'Wfinal',\n",
        "        shape=[height_final * width_final * channels_final, classes],\n",
        "        initializer=layers.xavier_initializer(),\n",
        "        trainable=True)\n",
        "    b_final = tf.get_variable(\n",
        "        'bfinal',\n",
        "        shape=[1, classes],\n",
        "        initializer=layers.xavier_initializer(),\n",
        "        trainable=True)\n",
        "\n",
        "    Ws.append(W_final)\n",
        "    bs.append(b_final)\n",
        "    net_hooks['weights']['Ws'] = Ws\n",
        "    net_hooks['weights']['bs'] = bs\n",
        "\n",
        "    a = tf.matmul(\n",
        "        tf.reshape(a, [-1, height_final * width_final * channels_final]),\n",
        "        W_final) + b_final\n",
        "    y = a\n",
        "    net_hooks['train_output']['y'] = y\n",
        "    y_ph_onehot = tf.one_hot(y_ph_in, classes, dtype=tf.int32)\n",
        "\n",
        "    # Weights and biases for regularization.\n",
        "    net_hooks['train_hook']['L2_reg_Ws'] = tf.reduce_sum(\n",
        "        [tf.reduce_sum(W_now**2.0) for W_now in Ws])\n",
        "    net_hooks['train_hook']['L2_reg_bs'] = tf.reduce_sum(\n",
        "        [tf.reduce_sum(b_now**2.0) for b_now in bs])\n",
        "\n",
        "    loss = tf.nn.softmax_cross_entropy_with_logits(labels=y_ph_onehot, logits=y)\n",
        "    net_hooks['train_hook']['loss'] = tf.reduce_mean(loss)\n",
        "  \n",
        "  # Defining all weights and biases as placeholders so that inference can be\n",
        "  # performed at given weight values.\n",
        "  with tf.variable_scope('eval'):\n",
        "    \n",
        "    Ws_ph = []\n",
        "    for W in Ws:\n",
        "      W_ph_now = tf.placeholder(tf.float32, W.get_shape())\n",
        "      Ws_ph.append(W_ph_now)\n",
        "    bs_ph = []\n",
        "    for b in bs:\n",
        "      b_ph_now = tf.placeholder(tf.float32, b.get_shape())\n",
        "      bs_ph.append(b_ph_now)\n",
        "\n",
        "    a = X_ph_in\n",
        "\n",
        "    for i in range(len(filter_sizes)):\n",
        "      _, _, _, Cnow = a.get_shape().as_list()\n",
        "\n",
        "      h = tf.nn.conv2d(\n",
        "          a, Ws_ph[i], strides=[1, 1, 1, 1], padding='SAME') + bs_ph[i]\n",
        "      h = tf.nn.dropout(h, rate=dropout_rate_ph_in)\n",
        "\n",
        "      h = tf.nn.max_pool(\n",
        "          h,\n",
        "          ksize=[1, pools[i], pools[i], 1],\n",
        "          strides=[1, pools[i], pools[i], 1],\n",
        "          padding='SAME')\n",
        "\n",
        "      if i \u003c len(filter_sizes) - 1:\n",
        "        a = f_nonlin(h)\n",
        "      else:\n",
        "        a = h\n",
        "      _, H_out, W_out, C_out = a.get_shape().as_list()\n",
        "    _, height_final, width_final, channels_final = a.get_shape().as_list()\n",
        "\n",
        "    last_layer = tf.reshape(a,\n",
        "                            [-1, height_final * width_final * channels_final])\n",
        "    net_hooks['train_hook']['last_layer'] = last_layer\n",
        "\n",
        "    a = tf.matmul(last_layer, Ws_ph[-1]) + bs_ph[-1]\n",
        "    y_eval = a\n",
        "    net_hooks['eval_output']['pred_eval'] = tf.nn.softmax(y_eval, axis=-1)\n",
        "    y_ph_onehot = tf.one_hot(y_ph_in, classes, dtype=tf.int32)\n",
        "    loss_eval = tf.nn.softmax_cross_entropy_with_logits(\n",
        "        labels=y_ph_onehot, logits=y_eval)\n",
        "    net_hooks['eval_output']['loss_eval'] = tf.reduce_mean(loss_eval)\n",
        "    net_hooks['eval_output']['y_eval'] = y_eval\n",
        "    net_hooks['placeholder']['Ws_ph'] = Ws_ph\n",
        "    net_hooks['placeholder']['bs_ph'] = bs_ph\n",
        "  return net_hooks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xebqAIwmq9ia"
      },
      "source": [
        "# Build model graph"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "height": 35
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 971,
          "status": "ok",
          "timestamp": 1596259216214,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "-WWIxPS1q5B6",
        "outputId": "7098c352-b897-434b-fc79-1972dbf29243"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of free parameters=971146\n"
          ]
        }
      ],
      "source": [
        "tf.reset_default_graph()\n",
        "\n",
        "X_ph = tf.placeholder(tf.float32, [None, Hn, Wn, Cn])\n",
        "y_ph = tf.placeholder(tf.int32, [None])\n",
        "lr_ph = tf.placeholder(tf.float32)\n",
        "dropout_rate_ph = tf.placeholder(tf.float32, [])\n",
        "L2_reg_constant_ph = tf.placeholder(tf.float32, [])\n",
        "\n",
        "architecture = 'CNN'\n",
        "\n",
        "if architecture == 'CNN':\n",
        "  with tf.variable_scope('to_get_shape'):\n",
        "    # Medium CNN\n",
        "    filter_sizes = (3, 3, 3, 3)\n",
        "    pools = (2, 2, 2, 2)\n",
        "    channels = (64, 128, 256, 256)\n",
        "\n",
        "    dummy_cnn_hooks = multilayer_CNN(\n",
        "        X_ph,\n",
        "        y_ph,\n",
        "        dropout_rate_ph_in=dropout_rate_ph,\n",
        "        filter_sizes=filter_sizes,\n",
        "        pools=pools,\n",
        "        channels=channels,\n",
        "        classes=classes)\n",
        "  # Number of parameters.\n",
        "  params = dummy_cnn_hooks['weights']['Ws'] + dummy_cnn_hooks['weights']['bs']\n",
        "  flat_params = tf.concat([tf.reshape(v, [-1]) for v in params], axis=0)\n",
        "  number_of_params = flat_params.get_shape().as_list()[0]\n",
        "  print('Number of free parameters=' + str(number_of_params))\n",
        "\n",
        "  cnn_hooks = multilayer_CNN(\n",
        "      X_ph,\n",
        "      y_ph,\n",
        "      dropout_rate_ph_in=dropout_rate_ph,\n",
        "      filter_sizes=filter_sizes,\n",
        "      pools=pools,\n",
        "      channels=channels,\n",
        "      classes=classes)\n",
        "  y = cnn_hooks['train_output']['y']\n",
        "  loss = cnn_hooks['train_hook']['loss']\n",
        "  y_eval = cnn_hooks['eval_output']['y_eval']\n",
        "  pred_eval = cnn_hooks['eval_output']['pred_eval']\n",
        "  loss_eval = cnn_hooks['eval_output']['loss_eval']\n",
        "  Ws_ph = cnn_hooks['placeholder']['Ws_ph']\n",
        "  bs_ph = cnn_hooks['placeholder']['bs_ph']\n",
        "  Ws = cnn_hooks['weights']['Ws']\n",
        "  bs = cnn_hooks['weights']['bs']\n",
        "  L2_reg_Ws = cnn_hooks['train_hook']['L2_reg_Ws']\n",
        "  L2_reg_bs = cnn_hooks['train_hook']['L2_reg_bs']\n",
        "\n",
        "loss_to_be_optimized = loss + L2_reg_constant_ph * (L2_reg_Ws + L2_reg_bs)\n",
        "\n",
        "train_step = tf.train.AdamOptimizer(lr_ph).minimize(loss_to_be_optimized)\n",
        "\n",
        "correct_prediction = tf.equal(tf.cast(tf.argmax(y, 1), dtype=tf.int32), y_ph)\n",
        "accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
        "\n",
        "correct_prediction_eval = tf.equal(\n",
        "    tf.cast(tf.argmax(y_eval, 1), dtype=tf.int32), y_ph)\n",
        "accuracy_eval = tf.reduce_mean(tf.cast(correct_prediction_eval, tf.float32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8PHJJ59SuMAo"
      },
      "source": [
        "## Helper Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eL52WPi6uPDD"
      },
      "outputs": [],
      "source": [
        "# Helper functions to flatten weights to a vector, and reform a flattened vector\n",
        "# to weights.\n",
        "def flatten(Ws1, bs1):\n",
        "  lists_now = []\n",
        "  for W_now in Ws1:\n",
        "    lists_now.append(W_now.reshape([-1]))\n",
        "  for b_now in bs1:\n",
        "    lists_now.append(b_now.reshape([-1]))\n",
        "  return np.concatenate(lists_now, axis=0)\n",
        "\n",
        "\n",
        "def get_flat_name(Ws_tf, bs_tf):\n",
        "  names = []\n",
        "  for W in Ws_tf:\n",
        "    names = names + [W.name] * np.prod(W.get_shape().as_list())\n",
        "  for b in bs_tf:\n",
        "    names = names + [b.name] * np.prod(b.get_shape().as_list())\n",
        "  return names\n",
        "\n",
        "\n",
        "def reform(flat1):\n",
        "  sofar = 0\n",
        "  Ws_out_now = []\n",
        "  bs_out_now = []\n",
        "  for W in Ws:\n",
        "    shape_now = W.get_shape().as_list()\n",
        "    size_now = np.prod(shape_now)\n",
        "    elements = flat1[sofar:sofar + size_now]\n",
        "    sofar = sofar + size_now\n",
        "    Ws_out_now.append(np.array(elements).reshape(shape_now))\n",
        "  for b in bs:\n",
        "    shape_now = b.get_shape().as_list()\n",
        "    size_now = np.prod(shape_now)\n",
        "    elements = flat1[sofar:sofar + size_now]\n",
        "    sofar = sofar + size_now\n",
        "    bs_out_now.append(np.array(elements).reshape(shape_now))\n",
        "  return Ws_out_now, bs_out_now\n",
        "\n",
        "\n",
        "# Brier score for evaluating uncertainty performance.\n",
        "def brier_scores(labels, probs=None, logits=None):\n",
        "  \"\"\"Compute elementwise Brier score.\n",
        "\n",
        "  Args:\n",
        "    labels: Tensor of integer labels shape [N1, N2, ...]\n",
        "    probs: Tensor of categorical probabilities of shape [N1, N2, ..., M].\n",
        "    logits: If `probs` is None, class probabilities are computed as a softmax\n",
        "      over these logits, otherwise, this argument is ignored.\n",
        "\n",
        "  Returns:\n",
        "    Tensor of shape [N1, N2, ...] consisting of Brier score contribution from\n",
        "    each element. The full-dataset Brier score is an average of these values.\n",
        "  \"\"\"\n",
        "  assert (probs is None) != (\n",
        "      logits is None), \"Exactly one of probs and logits should be None.\"\n",
        "  if probs is None:\n",
        "    probs = softmax(logits, axis=-1)\n",
        "  nlabels = probs.shape[-1]\n",
        "  flat_probs = probs.reshape([-1, nlabels])\n",
        "  flat_labels = labels.reshape([len(flat_probs)])\n",
        "\n",
        "  plabel = flat_probs[np.arange(len(flat_labels)), flat_labels]\n",
        "  out = np.square(flat_probs).sum(axis=-1) - 2 * plabel\n",
        "  return out.reshape(labels.shape) + 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "W2OrJv2tvdLv"
      },
      "source": [
        "# The root training loop (Independent solutions)\n",
        "This section could be time-consuming depending on the size of 'points_to_collect' ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "T6T6KJVWvNVP"
      },
      "outputs": [],
      "source": [
        "# Choosing a subset of the train data for faster eval.\n",
        "N_train_subset = N_val\n",
        "train_chosen_indices = np.random.choice(\n",
        "    range(N_train), N_train_subset, replace=False)\n",
        "X_train_subset = X_train[train_chosen_indices]\n",
        "y_train_subset = y_train[train_chosen_indices]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Ls0ziO7ivQo4"
      },
      "outputs": [],
      "source": [
        "# Number of independent solutions to collect.\n",
        "points_to_collect = 5\n",
        "# Number of ensemble models to use.\n",
        "max_ens_size = 4\n",
        "# Number of full training trajectory to collect for analysis such as T-SNE.\n",
        "num_trajectory_record = 3\n",
        "# Collect checkpoints along trajectory for subspace sampling or SWA.\n",
        "collect_solution_after_epoch = 30\n",
        "\n",
        "# Training hyperparams for each of the runs.\n",
        "epochs = 40\n",
        "batch_size = 128\n",
        "dropout = 0.1\n",
        "L2_constant = 0.0\n",
        "\n",
        "learning_rate = 1.6e-3\n",
        "# Defining the LR schedule\n",
        "lr_halving_epoch = 10\n",
        "lr_halving_maxcount = 1000\n",
        "lr_halving_power = 2.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Fw6WhPz-slnf"
      },
      "outputs": [],
      "source": [
        "print(\"train_step\", train_step)\n",
        "\n",
        "Ws_many = []\n",
        "bs_many = []\n",
        "\n",
        "losses_many = []\n",
        "accs_many = []\n",
        "\n",
        "# Collecting last epochs from each trajectory.\n",
        "Ws_by_epochs_many = [[] for _ in range(points_to_collect)]\n",
        "bs_by_epochs_many = [[] for _ in range(points_to_collect)]\n",
        "\n",
        "# Collecting whole trajectory.\n",
        "Ws_trajectory = [[] for _ in range(num_trajectory_record)]\n",
        "bs_trajectory = [[] for _ in range(num_trajectory_record)]\n",
        "\n",
        "global_init_op = tf.global_variables_initializer()\n",
        "with tf.Session() as sess:\n",
        "  for point_id in range(points_to_collect):\n",
        "    print(\"Optimization \" + str(point_id) + \" with starting lr=\" +\n",
        "          str(learning_rate) + \" dropout rate=\" + str(dropout) + \" batch=\" +\n",
        "          str(batch_size) + \" L2const=\" + str(L2_constant))\n",
        "    sess.run(global_init_op)\n",
        "    for e in range(epochs):\n",
        "      iterations = int(np.floor(float(N_train) / float(batch_size)))\n",
        "\n",
        "      losses_train_list = []\n",
        "      accs_train_list = []\n",
        "\n",
        "      halvings_count_now = np.floor(float(e) / lr_halving_epoch)\n",
        "      halvings_to_be_used = np.min([halvings_count_now, lr_halving_maxcount])\n",
        "      learning_rate_now = learning_rate / (\n",
        "          (lr_halving_power)**halvings_to_be_used)\n",
        "\n",
        "      for it in range(iterations):\n",
        "        indices = np.random.choice(range(N_train), batch_size)\n",
        "        X_batch = X_train[indices]\n",
        "        y_batch = y_train[indices]\n",
        "\n",
        "        feed_dict = {\n",
        "            X_ph: X_batch,\n",
        "            y_ph: y_batch,\n",
        "            lr_ph: learning_rate_now,\n",
        "            dropout_rate_ph: dropout,\n",
        "            L2_reg_constant_ph: L2_constant,\n",
        "        }\n",
        "\n",
        "        variables = [train_step, loss, accuracy]\n",
        "\n",
        "        _, loss_out, acc_out = sess.run(variables, feed_dict=feed_dict)\n",
        "\n",
        "        losses_train_list.append(loss_out)\n",
        "        accs_train_list.append(acc_out)\n",
        "\n",
        "      # Evaluating current epoch.\n",
        "      feed_dict = {\n",
        "          X_ph: X_val,\n",
        "          y_ph: y_val,\n",
        "          dropout_rate_ph: 0.0,\n",
        "          L2_reg_constant_ph: 0.0,\n",
        "      }\n",
        "\n",
        "      loss_val_out, acc_val_out, Ws_opt_out_now, bs_opt_out_now = sess.run(\n",
        "          [loss, accuracy, Ws, bs], feed_dict=feed_dict)\n",
        "\n",
        "      feed_dict[X_ph] = X_train_subset\n",
        "      feed_dict[y_ph] = y_train_subset\n",
        "\n",
        "      loss_train_out, acc_train_out = sess.run([loss, accuracy],\n",
        "                                               feed_dict=feed_dict)\n",
        "\n",
        "      feed_dict[X_ph] = X_test\n",
        "      feed_dict[y_ph] = y_test\n",
        "\n",
        "      loss_test_out, acc_test_out = sess.run([loss, accuracy],\n",
        "                                             feed_dict=feed_dict)\n",
        "\n",
        "      print(\"e=\" + str(e) + \" train loss=\" + f\"{loss_train_out:.4f}\" +\n",
        "            \" train acc=\" + f\"{acc_train_out:.4f}\" + \" val loss=\" +\n",
        "            f\"{loss_val_out:.4f}\" + \" val acc=\" + f\"{acc_val_out:.4f}\" +\n",
        "            \" test loss=\" + f\"{loss_test_out:.4f}\" + \" test acc=\" +\n",
        "            f\"{acc_test_out:.4f}\")\n",
        "      if e \u003e= collect_solution_after_epoch:\n",
        "        Ws_by_epochs_many[point_id].append(Ws_opt_out_now)\n",
        "        bs_by_epochs_many[point_id].append(bs_opt_out_now)\n",
        "      if point_id \u003c num_trajectory_record:\n",
        "        Ws_trajectory[point_id].append(Ws_opt_out_now)\n",
        "        bs_trajectory[point_id].append(bs_opt_out_now)\n",
        "\n",
        "    # Saving model weights of the last checkpoint.\n",
        "    Ws_opt_out_now, bs_opt_out_now = sess.run([Ws, bs])\n",
        "\n",
        "    Ws_many.append(Ws_opt_out_now)\n",
        "    bs_many.append(bs_opt_out_now)\n",
        "    losses_many.append(loss_val_out)\n",
        "    accs_many.append(acc_val_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "KapGgJupd9SZ"
      },
      "source": [
        "###Collecting Validation Predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8SJntXp5U2Ir"
      },
      "outputs": [],
      "source": [
        "trajectory_preds_test = np.zeros(\n",
        "    (num_trajectory_record, epochs, N_val, classes))\n",
        "independent_preds_test = np.zeros((points_to_collect, N_val, classes))\n",
        "\n",
        "with tf.Session() as sess:\n",
        "  for id_now in range(points_to_collect):\n",
        "    for j, W in enumerate(Ws_many[id_now]):\n",
        "      feed_dict[Ws_ph[j]] = W\n",
        "    for j, b in enumerate(bs_many[id_now]):\n",
        "      feed_dict[bs_ph[j]] = b\n",
        "\n",
        "    feed_dict[dropout_rate_ph] = 0.0\n",
        "\n",
        "    feed_dict[X_ph] = X_val\n",
        "    feed_dict[y_ph] = y_val\n",
        "    pred_eval_out = sess.run(pred_eval, feed_dict=feed_dict)\n",
        "    independent_preds_test[id_now] = pred_eval_out\n",
        "\n",
        "  for id_now in range(num_trajectory_record):\n",
        "    for e in range(epochs):\n",
        "      for j, W in enumerate(Ws_trajectory[id_now][e]):\n",
        "        feed_dict[Ws_ph[j]] = W\n",
        "      for j, b in enumerate(bs_trajectory[id_now][e]):\n",
        "        feed_dict[bs_ph[j]] = b\n",
        "      pred_eval_out = sess.run(pred_eval, feed_dict=feed_dict)\n",
        "      trajectory_preds_test[id_now][e] = pred_eval_out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MSyiWg_W9b4F"
      },
      "source": [
        "#Cosine and Fraction of Disagreement"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wVnTmgAD9l9Y"
      },
      "outputs": [],
      "source": [
        "def cos_between(v1, v2):\n",
        "  \"\"\" Returns the angle in radians between vectors 'v1' and 'v2'\"\"\"\n",
        "  v1_u = v1 / np.linalg.norm(v1)\n",
        "  v2_u = v2 / np.linalg.norm(v2)\n",
        "  return np.dot(v1_u, v2_u)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FV_a7LZ6BROc"
      },
      "source": [
        "### Within Trajectory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HLjL_tcXGAZp"
      },
      "outputs": [],
      "source": [
        "trajectory_id = 0\n",
        "flat_p_list = []\n",
        "num_epochs = len(Ws_trajectory[trajectory_id])\n",
        "cos_matrix = np.zeros((num_epochs, num_epochs))\n",
        "\n",
        "for e in range(num_epochs):\n",
        "  flat_p_list.append(\n",
        "      flatten(Ws_trajectory[trajectory_id][e], bs_trajectory[trajectory_id][e]))\n",
        "for i in range(num_epochs):\n",
        "  for j in range(i, num_epochs):\n",
        "    cos_matrix[i][j] = cos_between(flat_p_list[i], flat_p_list[j])\n",
        "    cos_matrix[j][i] = cos_matrix[i][j]\n",
        "\n",
        "plt.imshow(cos_matrix, interpolation=\"nearest\", cmap=\"bwr\", origin=\"lower\")\n",
        "plt.colorbar()\n",
        "plt.grid(\"off\")\n",
        "\n",
        "title = \"Cosine Along Train Trajectory\"\n",
        "plt.title(title)\n",
        "\n",
        "plt.xlabel(\"Checkpoint id\")\n",
        "plt.ylabel(\"Checkpoint id\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gprmfsa2n8a_"
      },
      "outputs": [],
      "source": [
        "preds_now = trajectory_preds_test[trajectory_id]\n",
        "targets_now = y_val\n",
        "\n",
        "classes_predicted = np.argmax(preds_now, axis=-1)\n",
        "fractional_differences = np.mean(\n",
        "    classes_predicted.reshape([1, epochs, len(targets_now)]) !=\n",
        "    classes_predicted.reshape([epochs, 1, len(targets_now)]),\n",
        "    axis=-1)\n",
        "\n",
        "plt.imshow(\n",
        "    fractional_differences, interpolation=\"nearest\", cmap=\"bwr\", origin=\"lower\")\n",
        "plt.colorbar()\n",
        "plt.grid(\"off\")\n",
        "\n",
        "title = \"Disagreement Fraction Along Train Trajectory\"\n",
        "plt.title(title)\n",
        "\n",
        "plt.xlabel(\"Checkpoint id\")\n",
        "plt.ylabel(\"Checkpoint id\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-FEPT3_OBT-Y"
      },
      "source": [
        "###Between Independent Runs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Wta3MB2tH4jB"
      },
      "outputs": [],
      "source": [
        "flat_p_list = []\n",
        "cos_matrix = np.zeros((points_to_collect, points_to_collect))\n",
        "for i in range(points_to_collect):\n",
        "  flat_p_list.append(flatten(Ws_many[i], bs_many[i]))\n",
        "\n",
        "for i in range(points_to_collect):\n",
        "  for j in range(i, points_to_collect):\n",
        "    cos_matrix[i][j] = cos_between(flat_p_list[i], flat_p_list[j])\n",
        "    cos_matrix[j][i] = cos_matrix[i][j]\n",
        "\n",
        "plt.imshow(cos_matrix, cmap=\"bwr\", origin=\"lower\")\n",
        "plt.colorbar()\n",
        "plt.grid(\"off\")\n",
        "\n",
        "title = \"Cosine Between Independent Solutions\"\n",
        "plt.title(title)\n",
        "\n",
        "plt.xlabel(\"Independent Solution\")\n",
        "plt.ylabel(\"Independent Solution\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Cglsoi44tCnL"
      },
      "outputs": [],
      "source": [
        "preds_now = independent_preds_test\n",
        "targets_now = y_val\n",
        "\n",
        "classes_predicted = np.argmax(preds_now, axis=-1)\n",
        "fractional_differences = np.mean(\n",
        "    classes_predicted.reshape(\n",
        "        [1, points_to_collect, len(targets_now)]) != classes_predicted.reshape(\n",
        "            [points_to_collect, 1, len(targets_now)]),\n",
        "    axis=-1)\n",
        "\n",
        "plt.imshow(\n",
        "    fractional_differences, interpolation=\"nearest\", cmap=\"bwr\", origin=\"lower\")\n",
        "plt.colorbar()\n",
        "plt.grid(\"off\")\n",
        "\n",
        "title = \"Disagreement Fraction Btw Independent Solutions\"\n",
        "plt.title(title)\n",
        "\n",
        "plt.xlabel(\"Independent Solution\")\n",
        "plt.ylabel(\"Independent Solution\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8hgROgxil29I"
      },
      "source": [
        "#T-SNE Visualization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "W6aVhQ5Yl7A6"
      },
      "outputs": [],
      "source": [
        "from sklearn.manifold import TSNE\n",
        "\n",
        "reshaped_prediction = trajectory_preds_test.reshape([-1, N_val * classes])\n",
        "\n",
        "prediction_embed = TSNE(n_components=2).fit_transform(reshaped_prediction)\n",
        "traj_embed = prediction_embed.reshape([num_trajectory_record, epochs, 2])\n",
        "\n",
        "colors_list = [\"r\", \"b\", \"g\"]\n",
        "labels_list = [\"traj_{}\".format(i) for i in range(num_trajectory_record)]\n",
        "for i in range(num_trajectory_record):\n",
        "  plt.plot(\n",
        "      traj_embed[i, :, 0],\n",
        "      traj_embed[i, :, 1],\n",
        "      color=colors_list[i],\n",
        "      alpha=0.8,\n",
        "      linestyle=\"\",\n",
        "      marker=\"o\",\n",
        "      label=labels_list[i])\n",
        "  plt.plot(\n",
        "      traj_embed[i, :, 0],\n",
        "      traj_embed[i, :, 1],\n",
        "      color=colors_list[i],\n",
        "      alpha=0.3,\n",
        "      linestyle=\"-\",\n",
        "      marker=\"\")\n",
        "  plt.plot(\n",
        "      traj_embed[i, 0, 0],\n",
        "      traj_embed[i, 0, 1],\n",
        "      color=colors_list[i],\n",
        "      alpha=1.0,\n",
        "      linestyle=\"\",\n",
        "      marker=\"*\",\n",
        "      markersize=20)\n",
        "plt.legend(loc=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ewL6qSnHlNEy"
      },
      "source": [
        "#Effects of Ensemble + Subspace Sampling"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FN3eJMBc55hG"
      },
      "source": [
        "## Gaussian Sampling\n",
        "Sample from model weight space according to a Gaussian distribution formed by last epoch checkpoints along a training trajectory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "d5TWmKrc53kK"
      },
      "outputs": [],
      "source": [
        "from sklearn.decomposition import PCA\n",
        "\n",
        "\n",
        "def get_gaussian_sample(var_mean, var_std, scale=1.0):\n",
        "  var_sample = np.random.normal(var_mean, scale * var_std)\n",
        "  return var_sample\n",
        "\n",
        "\n",
        "def get_pca_gaussian_flat_sampling(pca, means, rank, scale=1.0):\n",
        "  standard_normals = np.random.normal(loc=0.0, scale=scale, size=(rank))\n",
        "  shifts = pca.inverse_transform(standard_normals)\n",
        "  return shifts + means\n",
        "\n",
        "\n",
        "def get_rand_norm_direction(shape):\n",
        "  random_direction = np.random.normal(loc=0.0, scale=1.0, size=shape)\n",
        "  random_direction_normed = random_direction / np.linalg.norm(random_direction)\n",
        "  return random_direction_normed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mtGq9Qkj7iEG"
      },
      "outputs": [],
      "source": [
        "# PCA rank.\n",
        "rank = 5\n",
        "num_sample = 30\n",
        "\n",
        "dial_gaussian_whole_Ws = [[] for _ in range(points_to_collect)]\n",
        "dial_gaussian_whole_bs = [[] for _ in range(points_to_collect)]\n",
        "\n",
        "pca_gaussian_whole_Ws = [[] for _ in range(points_to_collect)]\n",
        "pca_gaussian_whole_bs = [[] for _ in range(points_to_collect)]\n",
        "\n",
        "for mid in range(points_to_collect):\n",
        "  Ws_traj = Ws_by_epochs_many[mid]\n",
        "  bs_traj = bs_by_epochs_many[mid]\n",
        "\n",
        "  vs_list = []\n",
        "  for i in range(len(Ws_traj)):\n",
        "    vs_list.append(flatten(Ws_traj[i], bs_traj[i]))\n",
        "\n",
        "  vs_np = np.stack(vs_list, axis=0)\n",
        "\n",
        "  means = np.mean(vs_np, axis=0)\n",
        "  stds = np.std(vs_np, axis=0)\n",
        "  vs_np_centered = vs_np - means.reshape([1, -1])\n",
        "\n",
        "  pca = PCA(n_components=rank)\n",
        "  pca.fit(vs_np_centered)\n",
        "  for i in range(num_sample):\n",
        "    v_sample = get_gaussian_sample(means, stds, scale=1.0)\n",
        "    w_sample, b_sample = reform(v_sample)\n",
        "    dial_gaussian_whole_Ws[mid].append(w_sample)\n",
        "    dial_gaussian_whole_bs[mid].append(b_sample)\n",
        "\n",
        "    v_sample = get_pca_gaussian_flat_sampling(pca, means, rank, scale=1.0)\n",
        "    w_sample, b_sample = reform(v_sample)\n",
        "\n",
        "    pca_gaussian_whole_Ws[mid].append(w_sample)\n",
        "    pca_gaussian_whole_bs[mid].append(b_sample)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U_kvSzLKzBSW"
      },
      "source": [
        "## Random Sampling\n",
        "Randomly sample from model weight space around the final checkpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PGILtMX7zUHp"
      },
      "outputs": [],
      "source": [
        "# Random samples need to meet this accuracy threshold to be included.\n",
        "acc_threshold = 0.70\n",
        "\n",
        "rand_Ws = [[] for _ in range(points_to_collect)]\n",
        "rand_bs = [[] for _ in range(points_to_collect)]\n",
        "\n",
        "with tf.Session() as sess:\n",
        "  feed_dict = {\n",
        "    X_ph: X_val,\n",
        "    y_ph: y_val,\n",
        "    lr_ph: 0.0,\n",
        "    dropout_rate_ph: 0.0,\n",
        "    L2_reg_constant_ph: 0.0,\n",
        "    }\n",
        "  for mid in range(points_to_collect):\n",
        "    for i in range(num_sample):\n",
        "      vs = flatten(Ws_many[mid], bs_many[mid])  \n",
        "      for trial in range(5):\n",
        "        scale = 10*np.random.uniform()   \n",
        "        v_sample = vs + scale * get_rand_norm_direction(vs.shape)\n",
        "        w_sample, b_sample = reform(v_sample)\n",
        "\n",
        "        for j,W in enumerate(w_sample):\n",
        "          feed_dict[Ws_ph[j]] = w_sample[j]\n",
        "        for j,b in enumerate(b_sample):\n",
        "          feed_dict[bs_ph[j]] = b_sample[j]  \n",
        "      \n",
        "        val_acc = sess.run(accuracy_eval,feed_dict = feed_dict)\n",
        "        if val_acc \u003e= acc_threshold:\n",
        "          rand_Ws[mid].append(w_sample)\n",
        "          rand_bs[mid].append(b_sample)\n",
        "          print('Obtaining 1 rand sample at scale {} with validation acc {} at {}th try'.format(scale, val_acc,trial))\n",
        "          break\n",
        "        if trial ==4:\n",
        "          print('No luck -------------------')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NJ5_RCriP-p5"
      },
      "outputs": [],
      "source": [
        "# PCA low-rank approximation of the random samplings\n",
        "pca_gaussian_rand_Ws = [[] for _ in range(points_to_collect)]\n",
        "pca_gaussian_rand_bs = [[] for _ in range(points_to_collect)]\n",
        "\n",
        "for mid in range(points_to_collect):\n",
        "  Ws_traj = rand_Ws[mid]\n",
        "  bs_traj = rand_bs[mid]\n",
        "\n",
        "  vs_list = []\n",
        "  for i in range(len(Ws_traj)):\n",
        "    vs_list.append(flatten(Ws_traj[i], bs_traj[i]))\n",
        "\n",
        "  vs_np = np.stack(vs_list, axis=0)\n",
        "\n",
        "  means = np.mean(vs_np, axis=0)\n",
        "  stds = np.std(vs_np, axis=0)\n",
        "  vs_np_centered = vs_np - means.reshape([1, -1])\n",
        "\n",
        "  pca = PCA(n_components=rank)\n",
        "  pca.fit(vs_np_centered)\n",
        "  for i in range(num_sample):\n",
        "    v_sample = get_pca_gaussian_flat_sampling(pca, means, rank, scale=1.0)\n",
        "    w_sample, b_sample = reform(v_sample)\n",
        "\n",
        "    pca_gaussian_rand_Ws[mid].append(w_sample)\n",
        "    pca_gaussian_rand_bs[mid].append(b_sample)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nEQ96SeTdnkZ"
      },
      "source": [
        "## Collecting predictions from Original and Subspace"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gKnPAilsoSNB"
      },
      "outputs": [],
      "source": [
        "# Average a list of weights.\n",
        "def average_var(w_list):\n",
        "  avg = [[] for _ in w_list[0]]\n",
        "  for w_now in w_list:\n",
        "    for i, w in enumerate(w_now):\n",
        "      avg[i].append(w)\n",
        "\n",
        "  for i, v in enumerate(avg):\n",
        "    avg[i] = np.mean(np.stack(v, axis=0), axis=0)\n",
        "\n",
        "  return avg\n",
        "\n",
        "\n",
        "# Given a list of model weights, feed_dict and a session, returns the model\n",
        "# predictions as a list.\n",
        "def get_pred_list(Ws_list, bs_list, feed_dict, sess):\n",
        "  pred_list = []\n",
        "  for id in range(len(Ws_list)):\n",
        "    Ws_now = Ws_list[id]\n",
        "    bs_now = bs_list[id]\n",
        "    for j, W in enumerate(Ws):\n",
        "      feed_dict[Ws_ph[j]] = Ws_now[j]\n",
        "    for j, b in enumerate(bs):\n",
        "      feed_dict[bs_ph[j]] = bs_now[j]\n",
        "\n",
        "    pred_eval_out = sess.run(pred_eval, feed_dict=feed_dict)\n",
        "    pred_list.append(pred_eval_out)\n",
        "  return pred_list\n",
        "\n",
        "\n",
        "# Consider a list of subspaces, each has a list of sampled weights. This\n",
        "# function computes model predictions, ensembles the predictions within each\n",
        "# subspace, and returns the list of ensembled predictions.\n",
        "def get_subspace_pred_list(Ws_subspace_list, bs_subspace_list, feed_dict, sess):\n",
        "  subspace_pred = []\n",
        "  num_subspace = len(Ws_subspace_list)\n",
        "  for mid in range(num_subspace):\n",
        "    pred_list_now = get_pred_list(Ws_subspace_list[mid], bs_subspace_list[mid],\n",
        "                                  feed_dict, sess)\n",
        "    subspace_pred.append(np.mean(np.stack(pred_list_now, axis=0), axis=0))\n",
        "  return subspace_pred"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "xeAuHy2xQ39g"
      },
      "outputs": [],
      "source": [
        "# Returns a list of all possible k-subset from [1, ..., n]\n",
        "# Don't scale well for large n. Use with caution.\n",
        "def choose_k_from_n(n, k):\n",
        "  if k\u003en or k \u003c1 or n \u003c 1:\n",
        "    return []\n",
        "  if k == n :\n",
        "    return [list(range(1, n+1))]\n",
        "  if k == 1:\n",
        "    return [[i] for i in range(1, n+1)]\n",
        "  a = choose_k_from_n(n-1,k)\n",
        "  b = choose_k_from_n(n-1, k-1)\n",
        "  b_new = []\n",
        "  for g in b:\n",
        "    b_new.append(g+[n])\n",
        "  return a+b_new\n",
        "\n",
        "\n",
        "def get_acc_brier(y, pred, is_logits=False):\n",
        "  acc = np.mean(np.argmax(pred,axis=-1)==y)\n",
        "  if is_logits:\n",
        "    brier = brier_scores(y,logits=pred)\n",
        "  else:\n",
        "    brier = brier_scores(y,probs=pred)\n",
        "  return acc, np.mean(brier)\n",
        "\n",
        "\n",
        "# Given a list of model predictions, compute the accuracy and brier score for \n",
        "# each individual model as well as the ensemble of them.\n",
        "def get_all_models_metrics(pred_list,y_test,max_ens_size=5, ens_size_list=None):\n",
        "  acc_list = []\n",
        "  acc_list_ensemble = []\n",
        "  b_list = []\n",
        "  b_list_ensemble = []\n",
        "  num_models = len(pred_list)\n",
        "  for i in range(num_models):\n",
        "    acc,brier = get_acc_brier(y_test, pred_list[i])\n",
        "    acc_list.append(acc)\n",
        "    b_list.append(brier)\n",
        "\n",
        "  max_ens_size = np.min([max_ens_size, num_models])\n",
        "  if ens_size_list is None:\n",
        "    ens_size_list = range(1, max_ens_size+1)\n",
        "  for ens_size in ens_size_list:\n",
        "    # Pick all possible subset with size of ens_size from available models.\n",
        "    # Compute ensemble for each such subset.\n",
        "    ens_index_list = choose_k_from_n(num_models,ens_size)\n",
        "    ens_acc = []\n",
        "    ens_brier=[]\n",
        "    for ens_ind in ens_index_list:\n",
        "      ens_pred_list = []\n",
        "      for ind in ens_ind:\n",
        "        ens_pred_list.append(pred_list[ind-1])\n",
        "      ens_np = np.mean(np.stack(ens_pred_list,axis=0), axis=0)\n",
        "      acc,brier = get_acc_brier(y_test, ens_np)\n",
        "      ens_acc.append(acc)\n",
        "      ens_brier.append(brier)\n",
        "    acc_list_ensemble.append(np.mean(ens_acc))\n",
        "    b_list_ensemble.append(np.mean(ens_brier))\n",
        "  metrics = {'accuracy': {},\n",
        "             'brier':{},}\n",
        "  metrics['accuracy']['individual'] = acc_list\n",
        "  metrics['accuracy']['ensemble'] = acc_list_ensemble\n",
        "  metrics['brier']['individual'] = b_list\n",
        "  metrics['brier']['ensemble'] = b_list_ensemble\n",
        "  return metrics\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RR2UKvMcdpJm"
      },
      "outputs": [],
      "source": [
        "# Get predictions on test data.\n",
        "orig_pred = []\n",
        "wa_pred = []\n",
        "\n",
        "# Compute averaged weights.\n",
        "wa_Ws = [[] for _ in range(points_to_collect)]\n",
        "wa_bs = [[] for _ in range(points_to_collect)]\n",
        "for i in range(points_to_collect):\n",
        "  wa_Ws[i] = average_var(Ws_by_epochs_many[i])\n",
        "  wa_bs[i] = average_var(bs_by_epochs_many[i])\n",
        "\n",
        "with tf.Session() as sess:\n",
        "  feed_dict[X_ph] = X_test\n",
        "  feed_dict[y_ph] = y_test\n",
        "  orig_pred = get_pred_list(Ws_many, bs_many, feed_dict, sess)\n",
        "  wa_pred = get_pred_list(wa_Ws, wa_bs, feed_dict, sess)\n",
        "  diag_gaus_pred = get_subspace_pred_list(dial_gaussian_whole_Ws,\n",
        "                                          dial_gaussian_whole_bs, feed_dict,\n",
        "                                          sess)\n",
        "  pca_gaus_pred = get_subspace_pred_list(pca_gaussian_whole_Ws,\n",
        "                                         pca_gaussian_whole_bs, feed_dict, sess)\n",
        "  pca_rand_pred = get_subspace_pred_list(pca_gaussian_rand_Ws,\n",
        "                                         pca_gaussian_rand_bs, feed_dict, sess)\n",
        "  rand_pred = get_subspace_pred_list(rand_Ws, rand_bs, feed_dict, sess)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sCSsNdA06lm1"
      },
      "outputs": [],
      "source": [
        "max_ens_size = points_to_collect - 1\n",
        "\n",
        "orig_metrics = get_all_models_metrics(\n",
        "    orig_pred, y_test, max_ens_size=max_ens_size)\n",
        "wa_metrics = get_all_models_metrics(wa_pred, y_test, max_ens_size=max_ens_size)\n",
        "\n",
        "diag_metrics = get_all_models_metrics(\n",
        "    diag_gaus_pred, y_test, max_ens_size=max_ens_size)\n",
        "pca_metrics = get_all_models_metrics(\n",
        "    pca_gaus_pred, y_test, max_ens_size=max_ens_size)\n",
        "\n",
        "rand_metrics = get_all_models_metrics(\n",
        "    rand_pred, y_test, max_ens_size=max_ens_size)\n",
        "pca_rand_metrics = get_all_models_metrics(\n",
        "    pca_rand_pred, y_test, max_ens_size=max_ens_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yQtoJGx817vV"
      },
      "source": [
        "##Plot Metrics for Ensembles + Subspace"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "utAi4c29CvgO"
      },
      "outputs": [],
      "source": [
        "title = \"Ensemble ACC test\"\n",
        "\n",
        "plt.xlabel(\"Ensemble size\")\n",
        "plt.ylabel(\"Test Acc\")\n",
        "\n",
        "ensemble_sizes = np.asarray(range(max_ens_size)) + 1\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    orig_metrics[\"accuracy\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"probs ensembling\",\n",
        "    color=\"navy\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    [np.mean(orig_metrics[\"accuracy\"][\"individual\"])] * len(ensemble_sizes),\n",
        "    label=\"original\",\n",
        "    color=\"blue\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    [np.mean(diag_metrics[\"accuracy\"][\"individual\"])] * len(ensemble_sizes),\n",
        "    label=\"Diag\",\n",
        "    color=\"pink\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    [np.mean(pca_metrics[\"accuracy\"][\"individual\"])] * len(ensemble_sizes),\n",
        "    label=\"PCA\",\n",
        "    color=\"green\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    diag_metrics[\"accuracy\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"Diag+Ensemble\",\n",
        "    color=\"red\")\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    wa_metrics[\"accuracy\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"WA+ensembling\",\n",
        "    color=\"grey\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    pca_metrics[\"accuracy\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"PCA+Ensemble\",\n",
        "    color=\"green\")\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    rand_metrics[\"accuracy\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"Rand+Ensemble\",\n",
        "    color=\"yellow\")\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    pca_rand_metrics[\"accuracy\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"PCA Rand+Ensemble\",\n",
        "    color=\"m\")\n",
        "\n",
        "plt.legend()\n",
        "plt.xlim(1, max_ens_size)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8vLQoHba6eXt"
      },
      "outputs": [],
      "source": [
        "title = \"Ensemble Brier test\"\n",
        "\n",
        "plt.xlabel(\"Ensemble size\")\n",
        "plt.ylabel(\"Test Brier\")\n",
        "\n",
        "ensemble_sizes = np.asarray(range(max_ens_size)) + 1\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    orig_metrics[\"brier\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"probs ensembling\",\n",
        "    color=\"navy\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    [np.mean(orig_metrics[\"brier\"][\"individual\"])] * len(ensemble_sizes),\n",
        "    label=\"original\",\n",
        "    color=\"blue\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    [np.mean(diag_metrics[\"brier\"][\"individual\"])] * len(ensemble_sizes),\n",
        "    label=\"Diag\",\n",
        "    color=\"pink\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    [np.mean(pca_metrics[\"brier\"][\"individual\"])] * len(ensemble_sizes),\n",
        "    label=\"PCA\",\n",
        "    color=\"green\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    diag_metrics[\"brier\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"Diag Ensemble\",\n",
        "    color=\"red\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    pca_metrics[\"brier\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"PCA Ensemble\",\n",
        "    color=\"green\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    wa_metrics[\"brier\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"WA ensembling\",\n",
        "    color=\"grey\")\n",
        "\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    rand_metrics[\"brier\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"Rand Ensemble\",\n",
        "    color=\"yellow\")\n",
        "plt.plot(\n",
        "    ensemble_sizes,\n",
        "    pca_rand_metrics[\"brier\"][\"ensemble\"],\n",
        "    marker=\"s\",\n",
        "    label=\"PCA Rand Ensemble\",\n",
        "    color=\"m\")\n",
        "\n",
        "plt.xlim(1, max_ens_size)\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bGelr5_KBL-O"
      },
      "source": [
        "## Evaluating on Corrupted CIFAR10\n",
        "See [paper](https://arxiv.org/pdf/1906.02530.pdf) for data description.\n",
        "\n",
        "CAUTION: This section will be very time-consuming. If one wants to test the code quickly, consider only run a small portion of CIFAR10-C -- reducing 'ALL_CORRUPTIONS' to containing only one type, and 'intensity_range' to [0, 1]."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "K2jeOP8AERPn"
      },
      "outputs": [],
      "source": [
        "ALL_CORRUPTIONS = [\n",
        "    'gaussian_noise',\n",
        "    'shot_noise',\n",
        "    # 'impulse_noise',\n",
        "    # 'defocus_blur',\n",
        "    # 'frosted_glass_blur',\n",
        "    # 'motion_blur',\n",
        "    # 'zoom_blur',\n",
        "    # 'snow',\n",
        "    # 'frost',\n",
        "    # 'fog',\n",
        "    # 'brightness',\n",
        "    # 'contrast',\n",
        "    # 'elastic',\n",
        "    # 'pixelate',\n",
        "    # 'jpeg_compression',\n",
        "    # 'gaussian_blur',\n",
        "    # 'saturate',\n",
        "    # 'spatter',\n",
        "    # 'speckle_noise',\n",
        "]\n",
        "\n",
        "num_models_to_use = 5\n",
        "ens_size_list = [2, 4]\n",
        "\n",
        "model_names = ['original', 'wa', 'diag', 'pca', 'rand', 'pca_rand']\n",
        "model_names_plus_ens = []\n",
        "for i in ens_size_list:\n",
        "  model_names_plus_ens = model_names_plus_ens + [\n",
        "      name + '_ens_{}'.format(i) for name in model_names\n",
        "  ]\n",
        "\n",
        "all_model_names = model_names + model_names_plus_ens\n",
        "\n",
        "model_to_acc_every_level = {name: [] for name in all_model_names}\n",
        "model_to_brier_every_level = {name: [] for name in all_model_names}\n",
        "\n",
        "intensity_range = range(6)\n",
        "for level in intensity_range:\n",
        "  print('========= Level {} ======='.format(level))\n",
        "  model_to_acc_this_level = {name: [] for name in all_model_names}\n",
        "  model_to_brier_this_level = {name: [] for name in all_model_names}\n",
        "  if level == 0:\n",
        "    corruptions = ['no_corruption']\n",
        "  else:\n",
        "    corruptions = ALL_CORRUPTIONS\n",
        "  for corruption_type in corruptions:\n",
        "    if corruption_type is 'no_corruption':\n",
        "      ds_test = tfds.load(name='cifar10', split=['test'], batch_size=-1)\n",
        "    else:\n",
        "      # Load corrupted data.\n",
        "      corruption_config_name = corruption_type + '_{}'.format(level)\n",
        "      ds_test = tfds.load(\n",
        "          name='cifar10_corrupted',\n",
        "          split=['test'],\n",
        "          builder_kwargs={'config': corruption_config_name},\n",
        "          batch_size=-1)\n",
        "    numpy_ds = tfds.as_numpy(ds_test)\n",
        "    x_test, y_test = numpy_ds[0]['image'], numpy_ds[0]['label']\n",
        "    x_test = x_test.reshape([-1, Hn, Wn, Cn])\n",
        "    y_test = y_test.reshape([-1])\n",
        "    N_test = len(y_test)\n",
        "    # Run inference\n",
        "    with tf.Session() as sess:\n",
        "      feed_dict[X_ph] = x_test\n",
        "      feed_dict[y_ph] = y_test\n",
        "      #Get predictions\n",
        "      for name in model_names:\n",
        "        if name is 'original':\n",
        "          pred_list = get_pred_list(Ws_many[0:num_models_to_use],\n",
        "                                    bs_many[0:num_models_to_use], feed_dict,\n",
        "                                    sess)\n",
        "        elif name is 'wa':\n",
        "          # Weight Averaged\n",
        "          pred_list = get_pred_list(wa_Ws[0:num_models_to_use],\n",
        "                                    wa_bs[0:num_models_to_use], feed_dict, sess)\n",
        "        elif name is 'diag':\n",
        "          # Subspacesampling\n",
        "          pred_list = get_subspace_pred_list(\n",
        "              dial_gaussian_whole_Ws[0:num_models_to_use],\n",
        "              dial_gaussian_whole_bs[0:num_models_to_use], feed_dict, sess)\n",
        "        elif name is 'pca':\n",
        "          pred_list = get_subspace_pred_list(\n",
        "              pca_gaussian_whole_Ws[0:num_models_to_use],\n",
        "              pca_gaussian_whole_bs[0:num_models_to_use], feed_dict, sess)\n",
        "        elif name is 'pca_rand':\n",
        "          pred_list = get_subspace_pred_list(\n",
        "              pca_gaussian_rand_Ws[0:num_models_to_use],\n",
        "              pca_gaussian_rand_bs[0:num_models_to_use], feed_dict, sess)\n",
        "        elif name is 'rand':\n",
        "          pred_list = get_subspace_pred_list(rand_Ws[0:num_models_to_use],\n",
        "                                             rand_bs[0:num_models_to_use],\n",
        "                                             feed_dict, sess)\n",
        "        corrupt_metrics = get_all_models_metrics(\n",
        "            pred_list, y_test, ens_size_list=ens_size_list)\n",
        "        model_to_acc_this_level[name].append(\n",
        "            np.mean(corrupt_metrics['accuracy']['individual']))\n",
        "        model_to_brier_this_level[name].append(\n",
        "            np.mean(corrupt_metrics['brier']['individual']))\n",
        "        for i, ens_size in enumerate(ens_size_list):\n",
        "          model_to_acc_this_level[name + '_ens_{}'.format(ens_size)].append(\n",
        "              corrupt_metrics['accuracy']['ensemble'][i])\n",
        "          model_to_brier_this_level[name + '_ens_{}'.format(ens_size)].append(\n",
        "              corrupt_metrics['brier']['ensemble'][i])\n",
        "\n",
        "  for name in all_model_names:\n",
        "    model_to_acc_every_level[name].append(model_to_acc_this_level[name])\n",
        "    model_to_brier_every_level[name].append(model_to_brier_this_level[name])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rGS9GPU2mo5j"
      },
      "source": [
        "## Plot the metrics across corruption intensity"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "b5cAs29LkPYu"
      },
      "outputs": [],
      "source": [
        "ens_size_list = [2, 4]\n",
        "model_names = ['original', 'wa', 'diag', 'pca', 'rand', 'pca_rand']\n",
        "model_names_plus_ens = []\n",
        "for i in ens_size_list:\n",
        "  model_names_plus_ens = model_names_plus_ens + [\n",
        "      name + '_ens_{}'.format(i) for name in model_names\n",
        "  ]\n",
        "\n",
        "all_model_names = model_names + model_names_plus_ens\n",
        "\n",
        "# Model names used by Jensen Shannon plots.\n",
        "js_model_names = ['inde', 'traj', 'diag', 'pca', 'rand', 'pca_rand']\n",
        "\n",
        "model_name_to_color = {\n",
        "    'original': 'indianred',\n",
        "    'wa': 'dimgray',\n",
        "    'diag': 'gold',\n",
        "    'pca': 'blue',\n",
        "    'rand': 'mediumseagreen',\n",
        "    'pca_rand': 'fuchsia',\n",
        "    'inde': 'red',\n",
        "    'traj': 'grey'\n",
        "}\n",
        "\n",
        "for name in model_names:\n",
        "  for i in ens_size_list:\n",
        "    model_name_to_color[name + '_ens_{}'.format(i)] = model_name_to_color[name]\n",
        "\n",
        "model_name_to_label = {\n",
        "    'original': 'Original',\n",
        "    'wa': 'Weight Avg',\n",
        "    'diag': 'Diag Gaus',\n",
        "    'pca': 'PCA Gaus',\n",
        "    'rand': 'Random',\n",
        "    'pca_rand': 'PCA Random',\n",
        "    'inde': 'Independent',\n",
        "    'traj': 'Train Trajectory'\n",
        "}\n",
        "\n",
        "for name in model_names:\n",
        "  model_name_to_label[name + '_ens'] = model_name_to_label[name] + '+Ens'\n",
        "  for i in ens_size_list:\n",
        "    model_name_to_label[\n",
        "        name +\n",
        "        '_ens_{}'.format(i)] = model_name_to_label[name] + '+Ens {}'.format(i)\n",
        "\n",
        "# Line style, currently can only support 2 ensemble sizes. Beyond that, solid\n",
        "# line will be used for all.\n",
        "model_name_to_ls = {}\n",
        "for name in model_names:\n",
        "  model_name_to_ls[name] = 'dashed'\n",
        "  model_name_to_ls[name + '_ens_{}'.format(ens_size_list[0])] = 'dotted'\n",
        "  for i in range(1, len(ens_size_list)):\n",
        "    model_name_to_ls[name + '_ens_{}'.format(ens_size_list[i])] = 'solid'\n",
        "\n",
        "\n",
        "def plot_metric_over_corrupted_data(metric_name, intensity_range,\n",
        "                                    all_model_names, base_model_names,\n",
        "                                    color_map, ls_map):\n",
        "  plt.figure(figsize=(10, 9))\n",
        "\n",
        "  ylabel = {\n",
        "      'acc': 'Accuracy',\n",
        "      'brier': 'Brier',\n",
        "  }\n",
        "  plt.xlabel('Corruption Intensity', fontsize=16)\n",
        "  plt.ylabel(ylabel[metric_name], fontsize=16)\n",
        "  label_map = {'original': 'Single'}\n",
        "  for ens_size in ens_size_list:\n",
        "    label_map['original_ens_{}'.format(ens_size)] = '{}-Ensemble'.format(\n",
        "        ens_size)\n",
        "  model_to_mean = {name: [] for name in all_model_names}\n",
        "  for name in all_model_names:\n",
        "    if metric_name is 'acc':\n",
        "      model_to_metric_every_level = model_to_acc_every_level[name]\n",
        "    elif metric_name is 'brier':\n",
        "      model_to_metric_every_level = model_to_brier_every_level[name]\n",
        "    for t in intensity_range:\n",
        "      model_to_mean[name].append(np.mean(model_to_metric_every_level[t]))\n",
        "\n",
        "    if name in label_map:\n",
        "      label = label_map[name]\n",
        "    else:\n",
        "      label = None\n",
        "    plt.plot(\n",
        "        intensity_range,\n",
        "        model_to_mean[name],\n",
        "        marker='s',\n",
        "        label=label,\n",
        "        color=color_map[name],\n",
        "        ls=ls_map[name])\n",
        "\n",
        "  legend0 = plt.legend(loc=3, fontsize=12)\n",
        "\n",
        "  patch_list = []\n",
        "  for name in base_model_names:\n",
        "    patch_list.append(\n",
        "        mpatch.Patch(\n",
        "            color=model_name_to_color[name],\n",
        "            label=model_name_to_label[name]))\n",
        "  legend1 = plt.legend(handles=patch_list)\n",
        "  ax = plt.gca().add_artist(legend0)\n",
        "\n",
        "  plt.xlim(np.min(intensity_range), np.max(intensity_range))\n",
        "  plt.tight_layout()\n",
        "\n",
        "\n",
        "def plot_js_over_corrupted_data(intensity_range, all_model_names, color_map,\n",
        "                                label_map):\n",
        "  plt.figure(figsize=(10, 9))\n",
        "  plt.xlabel('Corruption Intensity', fontsize=16)\n",
        "  plt.ylabel('Jensen Shannon', fontsize=16)\n",
        "\n",
        "  model_to_mean = {name: [] for name in all_model_names}\n",
        "\n",
        "  for name in all_model_names:\n",
        "    model_to_metric_every_level = model_to_js_every_level[name]\n",
        "    for t in intensity_range:\n",
        "      model_to_mean[name].append(np.mean(model_to_metric_every_level[t]))\n",
        "    plt.plot(\n",
        "        intensity_range,\n",
        "        model_to_mean[name],\n",
        "        marker='s',\n",
        "        label=label_map[name],\n",
        "        color=color_map[name])\n",
        "\n",
        "  plt.xlim(np.min(intensity_range), np.max(intensity_range))\n",
        "  plt.tight_layout()\n",
        "  plt.legend()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cszvIYjZkHLi"
      },
      "outputs": [],
      "source": [
        "intensity_range = range(6)\n",
        "plot_metric_over_corrupted_data('acc', intensity_range, all_model_names,\n",
        "                                model_names, model_name_to_color,\n",
        "                                model_name_to_ls)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aLsVQu1SYzaX"
      },
      "outputs": [],
      "source": [
        "intensity_range = range(6)\n",
        "plot_metric_over_corrupted_data('brier', intensity_range, all_model_names,\n",
        "                                model_names, model_name_to_color,\n",
        "                                model_name_to_ls)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZwOzmfhn3Z4l"
      },
      "source": [
        "##Jensen Shannon"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "tqiAL-z73cPe"
      },
      "outputs": [],
      "source": [
        "def get_entropy(p):\n",
        "  return np.sum(-p * np.log(p + 1e-5), axis=1)\n",
        "\n",
        "\n",
        "# Compute the Jensen Shannon score given a list of model predictions.\n",
        "def get_jensen_shannon(pred_list, is_logit=False):\n",
        "  n = len(pred_list)\n",
        "  if is_logit:\n",
        "    p_list = []\n",
        "    for i in range(n):\n",
        "      p_list.append(softmax(pred_list[i], axis=1))\n",
        "  else:\n",
        "    p_list = pred_list\n",
        "  p_np = np.stack(p_list, axis=0)\n",
        "  ensemble_pred = np.mean(p_np, axis=0)\n",
        "  ensemble_entropy = np.mean(get_entropy(ensemble_pred))\n",
        "\n",
        "  mean_entropy = np.mean(np.sum(-p_np * np.log(p_np + 1e-5), axis=2))\n",
        "  return ensemble_entropy - mean_entropy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mCgvOt_gbRYG"
      },
      "source": [
        "### Eval on SVHN\n",
        "JS on OOD data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4WralaRlbVOW"
      },
      "outputs": [],
      "source": [
        "_, ds_test = tfds.load(\n",
        "    name=\"svhn_cropped\", split=[\"train\", \"test\"], batch_size=-1)\n",
        "\n",
        "numpy_ds = tfds.as_numpy(ds_test)\n",
        "X_test_raw, y_test_fine_raw = numpy_ds[\"image\"], numpy_ds[\"label\"]\n",
        "\n",
        "X_test_svhn = X_test_raw\n",
        "y_test_svhn = y_test_fine_raw\n",
        "\n",
        "subsample = np.random.choice(len(y_test_svhn), 10000, replace=False)\n",
        "X_svhn = X_test_svhn[subsample, :, :, :]\n",
        "y_svhn = y_test_svhn[subsample]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ICYLKUto1mrv"
      },
      "outputs": [],
      "source": [
        "traj_svhn_js = []\n",
        "pca_svhn_js = []\n",
        "diag_svhn_js = []\n",
        "rand_svhn_js = []\n",
        "pca_rand_svhn_js = []\n",
        "\n",
        "with tf.Session() as sess:\n",
        "  feed_dict[X_ph] = X_svhn\n",
        "  feed_dict[y_ph] = y_svhn\n",
        "  feed_dict[lr_ph] = 0.0\n",
        "  feed_dict[dropout_rate_ph] = 0.0\n",
        "  indi_svhn_pred = get_pred_list(Ws_many[0:10], bs_many[0:10], feed_dict, sess)\n",
        "  for traj_id in range(points_to_collect):\n",
        "    trajectory_svhn_pred = get_pred_list(Ws_by_epochs_many[traj_id],\n",
        "                                         bs_by_epochs_many[traj_id], feed_dict,\n",
        "                                         sess)\n",
        "    traj_svhn_js.append(get_jensen_shannon(trajectory_svhn_pred))\n",
        "\n",
        "    diag_gaus_whole_svhn_pred = get_pred_list(dial_gaussian_whole_Ws[traj_id],\n",
        "                                              dial_gaussian_whole_bs[traj_id],\n",
        "                                              feed_dict, sess)\n",
        "    diag_svhn_js.append(get_jensen_shannon(diag_gaus_whole_svhn_pred))\n",
        "\n",
        "    pca_gaus_whole_svhn_pred = get_pred_list(pca_gaussian_whole_Ws[traj_id],\n",
        "                                             pca_gaussian_whole_bs[traj_id],\n",
        "                                             feed_dict, sess)\n",
        "    pca_svhn_js.append(get_jensen_shannon(pca_gaus_whole_svhn_pred))\n",
        "\n",
        "    pca_rand_svhn_pred = get_pred_list(pca_gaussian_rand_Ws[traj_id],\n",
        "                                       pca_gaussian_rand_bs[traj_id], feed_dict,\n",
        "                                       sess)\n",
        "    pca_rand_svhn_js.append(get_jensen_shannon(pca_rand_svhn_pred))\n",
        "\n",
        "    rand_svhn_pred = get_pred_list(rand_Ws[traj_id], rand_bs[traj_id],\n",
        "                                   feed_dict, sess)\n",
        "    rand_svhn_js.append(get_jensen_shannon(rand_svhn_pred))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-1j-4_KhAZEg"
      },
      "outputs": [],
      "source": [
        "print(\"Jensen-Shannon for Independent run is {:.4f}\".format(\n",
        "    get_jensen_shannon(indi_svhn_pred)))\n",
        "print(\"Jensen-Shannon for within-trajectory is {:.4f}\".format(\n",
        "    np.mean(traj_svhn_js)))\n",
        "\n",
        "print(\"Jensen-Shannon for Diag Gaussian is {:.4f}\".format(\n",
        "    np.mean(diag_svhn_js)))\n",
        "print(\"Jensen-Shannon for PCA Gaussian is {:.4f}\".format(np.mean(pca_svhn_js)))\n",
        "\n",
        "print(\"Jensen-Shannon for Rand is {:.4f}\".format(np.mean(rand_svhn_js)))\n",
        "print(\"Jensen-Shannon for PCA Rand is {:.4f}\".format(np.mean(pca_rand_svhn_js)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FLCIw3GwQeaa"
      },
      "source": [
        "###Eval on Cifar-10-C\n",
        "\n",
        "CAUTION: This section will be very time-consuming. If one wants to test the code quickly, consider only run a small portion of CIFAR10-C -- reducing 'ALL_CORRUPTIONS' to containing only one type, and 'intensity_range' to [0, 1]."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UNM4VieiQjAp"
      },
      "outputs": [],
      "source": [
        "# Choose any training trajectory to study the JS of subspace samplings.\n",
        "traj_id = 3\n",
        "\n",
        "model_to_js_every_level = {name: [] for name in js_model_names}\n",
        "\n",
        "intensity_range = range(6)\n",
        "for level in intensity_range:\n",
        "  print(\"========= Level {} =======\".format(level))\n",
        "  model_to_js_this_level = {name: [] for name in js_model_names}\n",
        "  if level == 0:\n",
        "    corruptions = [\"no_corruption\"]\n",
        "  else:\n",
        "    corruptions = ALL_CORRUPTIONS\n",
        "  for corruption_type in corruptions:\n",
        "    if corruption_type is \"no_corruption\":\n",
        "      ds_test = tfds.load(name=\"cifar10\", split=[\"test\"], batch_size=-1)\n",
        "    else:\n",
        "      # Load corrupted data.\n",
        "      corruption_config_name = corruption_type + \"_{}\".format(level)\n",
        "      ds_test = tfds.load(\n",
        "          name=\"cifar10_corrupted\",\n",
        "          split=[\"test\"],\n",
        "          builder_kwargs={\"config\": corruption_config_name},\n",
        "          batch_size=-1)\n",
        "    numpy_ds = tfds.as_numpy(ds_test)\n",
        "    x_test, y_test = numpy_ds[0][\"image\"], numpy_ds[0][\"label\"]\n",
        "    x_test = x_test.reshape([-1, Hn, Wn, Cn])\n",
        "    y_test = y_test.reshape([-1])\n",
        "    N_test = len(y_test)\n",
        "    # Run inference\n",
        "    with tf.Session() as sess:\n",
        "      feed_dict[X_ph] = x_test\n",
        "      feed_dict[y_ph] = y_test\n",
        "      #Get predictions\n",
        "      for name in js_model_names:\n",
        "        if name is \"inde\":\n",
        "          pred_list = get_pred_list(Ws_many, bs_many, feed_dict,\n",
        "                                    sess)\n",
        "        elif name is \"traj\":\n",
        "          # Weight Averaged\n",
        "          pred_list = get_pred_list(Ws_by_epochs_many[traj_id],\n",
        "                                    bs_by_epochs_many[traj_id], feed_dict, sess)\n",
        "        elif name is \"diag\":\n",
        "          pred_list = get_pred_list(dial_gaussian_whole_Ws[traj_id],\n",
        "                                    dial_gaussian_whole_bs[traj_id], feed_dict,\n",
        "                                    sess)\n",
        "        elif name is \"pca\":\n",
        "          pred_list = get_pred_list(pca_gaussian_whole_Ws[traj_id],\n",
        "                                    pca_gaussian_whole_bs[traj_id], feed_dict,\n",
        "                                    sess)\n",
        "        elif name is \"pca_rand\":\n",
        "          pred_list = get_pred_list(pca_gaussian_rand_Ws[traj_id],\n",
        "                                    pca_gaussian_rand_bs[traj_id], feed_dict,\n",
        "                                    sess)\n",
        "        elif name is \"rand\":\n",
        "          pred_list = get_pred_list(rand_Ws[traj_id], rand_bs[traj_id],\n",
        "                                    feed_dict, sess)\n",
        "\n",
        "        model_to_js_this_level[name].append(get_jensen_shannon(pred_list))\n",
        "\n",
        "  for name in js_model_names:\n",
        "    model_to_js_every_level[name].append(model_to_js_this_level[name])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2mcbo-46d4s2"
      },
      "outputs": [],
      "source": [
        "plot_js_over_corrupted_data(intensity_range, js_model_names,\n",
        "                            model_name_to_color, model_name_to_label)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zCZUtnjKljz-"
      },
      "source": [
        "#Diversity vs Accuracy\n",
        "This section plots the Diversiy vs Accuracy plane for different subspace sampling methods, compared to independent solutions:\n",
        "\n",
        "\n",
        "*   Random\n",
        "*   Gaussian Diag\n",
        "*   Gaussian Low rank (PCA)\n",
        "*   Dropout\n",
        "\n",
        "Note that unlike the previous section \"Effects of Ensemble + Subspace Sampling\" where we only keep samples with good accuracy performance, in this section we include samples that sacrifices accuracy to explore the full spectrum of function diversity"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D3eKK8UPz-UB"
      },
      "outputs": [],
      "source": [
        "base_id = 0\n",
        "base_v = flatten(Ws_many[base_id], bs_many[base_id])\n",
        "\n",
        "\n",
        "def get_acc_and_diff(pred_class, base_pred, y_label):\n",
        "  diff = np.mean(pred_class != base_pred)\n",
        "  acc = np.mean(pred_class == y_label)\n",
        "  return acc, diff\n",
        "\n",
        "\n",
        "# Given a list of weights, compute their accuracy and their difference to the\n",
        "# base prediction.\n",
        "def get_acc_and_diff_from_weights(Ws_list, bs_list, feed_dict, base_pred):\n",
        "  assert len(Ws_list) == len(bs_list)\n",
        "  with tf.Session() as sess:\n",
        "    pred_list = get_pred_list(Ws_list, bs_list, feed_dict, sess)\n",
        "\n",
        "  acc_list = []\n",
        "  diff_list = []\n",
        "  for i in range(len(Ws_list)):\n",
        "    acc, diff = get_acc_and_diff(\n",
        "        np.argmax(pred_list[i], axis=-1), base_pred, feed_dict[y_ph])\n",
        "    acc_list.append(acc)\n",
        "    diff_list.append(diff)\n",
        "  return acc_list, diff_list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0oScJ4X09-F9"
      },
      "source": [
        "##Independent Solutions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TkwIdoOJvgz2"
      },
      "outputs": [],
      "source": [
        "feed_dict = {\n",
        "    X_ph: X_test,\n",
        "    y_ph: y_test,\n",
        "    lr_ph: 0.0,\n",
        "    dropout_rate_ph: 0.0,\n",
        "    L2_reg_constant_ph: 0.0,\n",
        "}\n",
        "\n",
        "with tf.Session() as sess:\n",
        "  orig_pred_list = get_pred_list(Ws_many, bs_many, feed_dict, sess)\n",
        "\n",
        "base_pred = np.argmax(orig_pred_list[base_id], axis=-1)\n",
        "\n",
        "independent_acc = []\n",
        "independent_diff = []\n",
        "for i in range(points_to_collect):\n",
        "  class_pred = np.argmax(orig_pred_list[i], axis=-1)\n",
        "  acc, diff = get_acc_and_diff(class_pred, base_pred, y_test)\n",
        "  independent_acc.append(acc)\n",
        "  independent_diff.append(diff)\n",
        "\n",
        "data_to_show = [\n",
        "    (\"independent optima\", \"red\", 300, \"*\", 1.0, independent_acc,\n",
        "     independent_diff),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cotIEVchK3Vh"
      },
      "outputs": [],
      "source": [
        "data_to_show = [\n",
        "    (\"independent optima\", \"red\", 300, \"*\", 1.0, independent_acc,\n",
        "     independent_diff),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hT1pKG9W-BTZ"
      },
      "source": [
        "##Gaussian Sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "czoWhgyGAFoC"
      },
      "outputs": [],
      "source": [
        "num_sample = 50\n",
        "rank = 5  # for PCA\n",
        "\n",
        "Ws_traj = Ws_by_epochs_many[base_id]\n",
        "bs_traj = bs_by_epochs_many[base_id]\n",
        "\n",
        "dial_gaussian_whole_Ws = []\n",
        "dial_gaussian_whole_bs = []\n",
        "\n",
        "pca_gaussian_whole_Ws = []\n",
        "pca_gaussian_whole_bs = []\n",
        "\n",
        "vs_list = []\n",
        "for i in range(len(Ws_traj)):\n",
        "  vs_list.append(flatten(Ws_traj[i], bs_traj[i]))\n",
        "\n",
        "vs_np = np.stack(vs_list, axis=0)\n",
        "\n",
        "means = np.mean(vs_np, axis=0)\n",
        "stds = np.std(vs_np, axis=0)\n",
        "vs_np_centered = vs_np - means.reshape([1, -1])\n",
        "\n",
        "pca = PCA(n_components=rank)\n",
        "pca.fit(vs_np_centered)\n",
        "for i in range(num_sample):\n",
        "  scale = np.random.uniform()\n",
        "  # One can adjust the constant in front of scale to get a fuller range in\n",
        "  # diversity-acc plane.\n",
        "  v_sample = get_gaussian_sample(means, stds, scale=10.0 * scale)\n",
        "  w_sample, b_sample = reform(v_sample)\n",
        "  dial_gaussian_whole_Ws.append(w_sample)\n",
        "  dial_gaussian_whole_bs.append(b_sample)\n",
        "  # One can adjust the scale value to get a fuller range in diversity-acc plane.\n",
        "  v_sample = get_pca_gaussian_flat_sampling(pca, means, rank, scale=6.0)\n",
        "  w_sample, b_sample = reform(v_sample)\n",
        "\n",
        "  pca_gaussian_whole_Ws.append(w_sample)\n",
        "  pca_gaussian_whole_bs.append(b_sample)\n",
        "\n",
        "pca_acc, pca_diff = get_acc_and_diff_from_weights(pca_gaussian_whole_Ws,\n",
        "                                                  pca_gaussian_whole_bs,\n",
        "                                                  feed_dict, base_pred)\n",
        "diag_acc, diag_diff = get_acc_and_diff_from_weights(dial_gaussian_whole_Ws,\n",
        "                                                    dial_gaussian_whole_bs,\n",
        "                                                    feed_dict, base_pred)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mfnrhHSa5Xiw"
      },
      "outputs": [],
      "source": [
        "data_to_show = data_to_show + [\n",
        "    (\"diag gaussian\", \"purple\", 10, \"o\", 0.6, diag_acc, diag_diff),\n",
        "    (\"pca gaussian\", \"fuchsia\", 10, \"o\", 0.3, pca_acc, pca_diff),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jn-g6-GZZ5tc"
      },
      "source": [
        "##Random Sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uJpRJaDLaFRa"
      },
      "outputs": [],
      "source": [
        "rand_Ws = []\n",
        "rand_bs = []\n",
        "for i in range(num_sample):\n",
        "  # One can adjust the constant in front of scale to get a fuller range in\n",
        "  # diversity-acc plane. \n",
        "  scale = 70.0 * np.random.uniform()\n",
        "  v_sample = base_v + scale * get_rand_norm_direction(base_v.shape)\n",
        "  w_sample, b_sample = reform(v_sample)\n",
        "  rand_Ws.append(w_sample)\n",
        "  rand_bs.append(b_sample)\n",
        "\n",
        "rand_acc, rand_diff = get_acc_and_diff_from_weights(rand_Ws, rand_bs, feed_dict,\n",
        "                                                    base_pred)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "h1my3OECbzEi"
      },
      "outputs": [],
      "source": [
        "data_to_show = data_to_show + [\n",
        "    (\"random subspace\", \"blue\", 3, \"o\", 0.3, rand_acc, rand_diff),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JJWIEJZBZ9Jf"
      },
      "source": [
        "##Dropout Sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "MBiCleCgb8r1"
      },
      "outputs": [],
      "source": [
        "def apply_dropout_to_array(anchor_array, dropout_to_use):\n",
        "  shape_now = anchor_array.shape\n",
        "  random_mask = np.random.rand(*shape_now)\n",
        "  mask = (random_mask \u003c (1.0 - dropout_to_use))\n",
        "  array_dropped = (anchor_array * mask) / (1.0 - dropout_to_use)\n",
        "  return array_dropped\n",
        "\n",
        "\n",
        "dropout_Ws = []\n",
        "dropout_bs = []\n",
        "for i in range(num_sample):\n",
        "  # One can adjust the constant to get a fuller range in diversity-acc plane.\n",
        "  dropout = 0.1 * np.random.uniform()\n",
        "  v_sample = apply_dropout_to_array(base_v, dropout)\n",
        "  w_sample, b_sample = reform(v_sample)\n",
        "  dropout_Ws.append(w_sample)\n",
        "  dropout_bs.append(b_sample)\n",
        "\n",
        "dropout_acc, dropout_diff = get_acc_and_diff_from_weights(\n",
        "    dropout_Ws, dropout_bs, feed_dict, base_pred)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "E5qdldWHb9c4"
      },
      "outputs": [],
      "source": [
        "data_to_show = data_to_show + [\n",
        "    (\"dropout subspace\", \"orange\", 10, \"o\", 1.0, dropout_acc, dropout_diff),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VwNL8PPM7DrM"
      },
      "source": [
        "##Plotting Figure"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AqhcI78L52hM"
      },
      "outputs": [],
      "source": [
        "# Functions for computing the analytic limit curves (see paper).\n",
        "\n",
        "def perturbed_reference_analytic(desired_accuracy, reference_accuracy, classes):\n",
        "  return (reference_accuracy - desired_accuracy) / (\n",
        "      reference_accuracy + (reference_accuracy - 1.0) / (classes - 1.0))\n",
        "\n",
        "\n",
        "def random_average_case(desired_accuracy, reference_accuracy, classes):\n",
        "  part1 = reference_accuracy * (1.0 - desired_accuracy)\n",
        "  part2 = desired_accuracy * (1.0 - reference_accuracy)\n",
        "  part3 = (1.0 - reference_accuracy) * (1.0 -\n",
        "                                        desired_accuracy) * (classes - 2.0) / (\n",
        "                                            classes - 1.0)\n",
        "  return part1 + part2 + part3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DjOAuW3F56qu"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(7, 5))\n",
        "accs_fit_random = np.linspace(0.1, 0.71, 100)\n",
        "diffs_fit_random = random_average_case(accs_fit_random, 0.71, 10)\n",
        "plt.plot(\n",
        "    accs_fit_random,\n",
        "    diffs_fit_random / (1.0 - accs_fit_random),\n",
        "    color=\"black\",\n",
        "    linestyle=\"-.\",\n",
        "    label=\"Upper limit\")\n",
        "\n",
        "# Analytic limits curve.\n",
        "accs_fit_perturbed = np.linspace(0.1, 0.71, 100)\n",
        "diffs_fit_perturbed = perturbed_reference_analytic(accs_fit_perturbed, 0.71, 10)\n",
        "plt.plot(\n",
        "    accs_fit_perturbed,\n",
        "    diffs_fit_perturbed / (1.0 - accs_fit_perturbed),\n",
        "    color=\"black\",\n",
        "    linestyle=\"--\",\n",
        "    label=\"Lower limit\")\n",
        "\n",
        "for (name_now, color_now, size_now, marker_now, alpha_now, accs_now,\n",
        "     diffs_now) in data_to_show:\n",
        "\n",
        "  metric_now = np.asarray(diffs_now) / (1.0 - np.asarray(accs_now))\n",
        "\n",
        "  if name_now != \"independent optima\":\n",
        "    plt.scatter(\n",
        "        accs_now,\n",
        "        metric_now,\n",
        "        color=color_now,\n",
        "        s=size_now,\n",
        "        marker=marker_now,\n",
        "        alpha=alpha_now)\n",
        "    plt.scatter([], [],\n",
        "                color=color_now,\n",
        "                label=name_now,\n",
        "                s=10,\n",
        "                marker=marker_now,\n",
        "                alpha=1.0)\n",
        "  else:\n",
        "    plt.scatter(\n",
        "        accs_now,\n",
        "        metric_now,\n",
        "        color=color_now,\n",
        "        s=size_now,\n",
        "        marker=marker_now,\n",
        "        alpha=alpha_now,\n",
        "        label=name_now)\n",
        "\n",
        "  if name_now == \"independent optima\":\n",
        "    base_star_color = \"green\"\n",
        "    plt.scatter([accs_now[base_id]], [metric_now[base_id]],\n",
        "                color=base_star_color,\n",
        "                label=\"baseline optimum\",\n",
        "                s=size_now,\n",
        "                marker=marker_now,\n",
        "                alpha=alpha_now)\n",
        "\n",
        "plt.xlabel(\"Validation accuracy\", fontsize=14)\n",
        "plt.ylabel(\"Fraction of labels changes / (1.0-accuracy)\", fontsize=14)\n",
        "\n",
        "plt.title(\"MediumCNN on Cifar10\", fontsize=18)\n",
        "plt.legend(loc=3, fancybox=True, framealpha=0.5, fontsize=16)\n",
        "plt.ylim([-0.1, 1.2])\n",
        "plt.xlim([0, 0.78])\n",
        "\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "8PHJJ59SuMAo"
      ],
      "name": "Release cifar10 medium cnn experiments.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}