# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + colab={} colab_type="code" id="WU--cjOphQKk"
  # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.

  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at

  #      http://www.apache.org/licenses/LICENSE-2.0

  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #  See the License for the specific language governing permissions and
  #  limitations under the License.

# + [markdown] colab_type="text" id="MGjUZ2Q55OWM"
# #Ensemble and Subspace Sampling Experiments on CIFAR10
#
#    <!-- Copyright 2020 <NAME>, <NAME>, <NAME>
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License. -->
#
#
#
# This notebook illustrates the CIFAR-10 experiments in the paper: 
#
# [Deep Ensembles: A Loss Landscape Perspective](https://arxiv.org/abs/1912.02757) by <NAME>, <NAME> and <NAME>
#
#
# These experiments investigate the effects of ensembling and variational Bayesian methods, please see the paper for more details.

# + [markdown] colab_type="text" id="U5sMeD_4pTe2"
# # Setting up

# + cellView="code" colab={} colab_type="code" id="bUXPnWQIEKJh"
# %tensorflow_version 1.x

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.contrib import layers
import tensorflow_datasets as tfds

from scipy.special import softmax
from matplotlib import patches as mpatch

# Plot Style
mpl.style.use('seaborn-colorblind')
mpl.rcParams.update({
    'font.size': 14,
    'lines.linewidth': 2,
    'figure.figsize': (6, 6 / 1.61)
})
mpl.rcParams['grid.color'] = 'k'
mpl.rcParams['grid.linestyle'] = ':'
mpl.rcParams['grid.linewidth'] = 0.5
mpl.rcParams['lines.markersize'] = 6
mpl.rcParams['lines.marker'] = None
mpl.rcParams['axes.grid'] = True

DEFAULT_FONTSIZE = 13
mpl.rcParams.update({
    'font.size': DEFAULT_FONTSIZE,
    'lines.linewidth': 2,
    'legend.fontsize': DEFAULT_FONTSIZE,
    'axes.labelsize': DEFAULT_FONTSIZE,
    'xtick.labelsize': DEFAULT_FONTSIZE,
    'ytick.labelsize': DEFAULT_FONTSIZE,
    'figure.figsize': (7, 7.0 / 1.4)
})


# + [markdown] colab_type="text" id="8_a0OnCQrkFu"
# #Getting the datasets

# + colab={} colab_type="code" id="AlG96jyJr8cW"
def give_me_data():
  print("Reading CIFAR10")
  cifar_data = {}
  N_val = 500

  # Construct a tf.data.Dataset
  ds_train, ds_test = tfds.load(
      name="cifar10", split=["train", "test"], batch_size=-1)

  numpy_train = tfds.as_numpy(ds_train)
  x_train_raw, y_train_raw = numpy_train["image"], numpy_train["label"]

  numpy_test = tfds.as_numpy(ds_test)
  x_test_raw, y_test_raw = numpy_test["image"], numpy_test["label"]

  N_train = x_train_raw.shape[0] - N_val

  X_train = x_train_raw[:N_train]
  y_train = y_train_raw[:N_train]
  X_val = x_train_raw[N_train:N_train + N_val]
  y_val = y_train_raw[N_train:N_train + N_val]
  X_test = x_test_raw
  y_test = y_test_raw

  Hn = 32
  Wn = 32
  Cn = 3
  cifar_data['Hn'] = Hn
  cifar_data['Wn'] = Wn
  cifar_data['Cn'] = Cn
  cifar_data['classes'] = 10

  cifar_data['X_train'] = X_train.reshape([-1, Hn, Wn, Cn])
  cifar_data['X_val'] = X_val.reshape([-1, Hn, Wn, Cn])
  cifar_data['X_test'] = X_test.reshape([-1, Hn, Wn, Cn])

  cifar_data['y_train'] = y_train.reshape([-1])
  cifar_data['y_val'] = y_val.reshape([-1])
  cifar_data['y_test'] = y_test.reshape([-1])
  return cifar_data


# + 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"
cifar_ds = give_me_data()
X_train = cifar_ds['X_train']
y_train = cifar_ds['y_train']
X_val = cifar_ds['X_val']
y_val = cifar_ds['y_val']
X_test = cifar_ds['X_test']
y_test = cifar_ds['y_test']
Hn = cifar_ds['Hn']
Wn = cifar_ds['Wn']
Cn = cifar_ds['Cn']
classes = cifar_ds['classes']

N_train = X_train.shape[0]
N_val = X_val.shape[0]
N_test = X_test.shape[0]


# + [markdown] colab_type="text" id="DsEjVGc_pqAL"
# # Defining a CNN

# + colab={} colab_type="code" id="fWU-qKjRvvKV"
# Builds a CNN network and returns various graph nodes for training and
# evaluation.
# In the 'eval' part of the graph, placeholders for network weights (Ws & bs)
# are created to facilitate inferencing at given weights (e.g. from subspace
# sampling).
def multilayer_CNN(X_ph_in,
                   y_ph_in,
                   dropout_rate_ph_in=None,
                   filter_sizes=(3, 3, 3, 3),
                   pools=(2, 2, 2, 2),
                   channels=(32, 64, 128, 256),
                   classes=10):
  net_hooks = {
      'weights': {},
      'placeholder': {},
      'train_hook': {},
      'train_output': {},
      'eval_output': {}
  }
  f_nonlin = tf.nn.relu
  with tf.variable_scope('train'):
    a = X_ph_in
    Ws = []
    bs = []
    for i in range(len(filter_sizes)):
      _, _, _, Cnow = a.get_shape().as_list()
      W = tf.get_variable(
          'Wconv' + str(i),
          shape=[filter_sizes[i], filter_sizes[i], Cnow, channels[i]],
          initializer=layers.xavier_initializer(),
          trainable=True)
      b = tf.get_variable(
          'bconv' + str(i),
          shape=[1, 1, channels[i]],
          initializer=layers.xavier_initializer(),
          trainable=True)
      Ws.append(W)
      bs.append(b)

      h = tf.nn.conv2d(a, W, strides=[1, 1, 1, 1], padding='SAME') + b
      h = tf.nn.dropout(h, rate=dropout_rate_ph_in)
      h = tf.nn.max_pool(
          h,
          ksize=[1, pools[i], pools[i], 1],
          strides=[1, pools[i], pools[i], 1],
          padding='SAME')
      if i < len(filter_sizes) - 1:
        a = f_nonlin(h)
      else:
        a = h
      _, H_out, W_out, C_out = a.get_shape().as_list()
    _, height_final, width_final, channels_final = a.get_shape().as_list()

    # Final fully connected layer.
    W_final = tf.get_variable(
        'Wfinal',
        shape=[height_final * width_final * channels_final, classes],
        initializer=layers.xavier_initializer(),
        trainable=True)
    b_final = tf.get_variable(
        'bfinal',
        shape=[1, classes],
        initializer=layers.xavier_initializer(),
        trainable=True)

    Ws.append(W_final)
    bs.append(b_final)
    net_hooks['weights']['Ws'] = Ws
    net_hooks['weights']['bs'] = bs

    a = tf.matmul(
        tf.reshape(a, [-1, height_final * width_final * channels_final]),
        W_final) + b_final
    y = a
    net_hooks['train_output']['y'] = y
    y_ph_onehot = tf.one_hot(y_ph_in, classes, dtype=tf.int32)

    # Weights and biases for regularization.
    net_hooks['train_hook']['L2_reg_Ws'] = tf.reduce_sum(
        [tf.reduce_sum(W_now**2.0) for W_now in Ws])
    net_hooks['train_hook']['L2_reg_bs'] = tf.reduce_sum(
        [tf.reduce_sum(b_now**2.0) for b_now in bs])

    loss = tf.nn.softmax_cross_entropy_with_logits(labels=y_ph_onehot, logits=y)
    net_hooks['train_hook']['loss'] = tf.reduce_mean(loss)
  
  # Defining all weights and biases as placeholders so that inference can be
  # performed at given weight values.
  with tf.variable_scope('eval'):
    
    Ws_ph = []
    for W in Ws:
      W_ph_now = tf.placeholder(tf.float32, W.get_shape())
      Ws_ph.append(W_ph_now)
    bs_ph = []
    for b in bs:
      b_ph_now = tf.placeholder(tf.float32, b.get_shape())
      bs_ph.append(b_ph_now)

    a = X_ph_in

    for i in range(len(filter_sizes)):
      _, _, _, Cnow = a.get_shape().as_list()

      h = tf.nn.conv2d(
          a, Ws_ph[i], strides=[1, 1, 1, 1], padding='SAME') + bs_ph[i]
      h = tf.nn.dropout(h, rate=dropout_rate_ph_in)

      h = tf.nn.max_pool(
          h,
          ksize=[1, pools[i], pools[i], 1],
          strides=[1, pools[i], pools[i], 1],
          padding='SAME')

      if i < len(filter_sizes) - 1:
        a = f_nonlin(h)
      else:
        a = h
      _, H_out, W_out, C_out = a.get_shape().as_list()
    _, height_final, width_final, channels_final = a.get_shape().as_list()

    last_layer = tf.reshape(a,
                            [-1, height_final * width_final * channels_final])
    net_hooks['train_hook']['last_layer'] = last_layer

    a = tf.matmul(last_layer, Ws_ph[-1]) + bs_ph[-1]
    y_eval = a
    net_hooks['eval_output']['pred_eval'] = tf.nn.softmax(y_eval, axis=-1)
    y_ph_onehot = tf.one_hot(y_ph_in, classes, dtype=tf.int32)
    loss_eval = tf.nn.softmax_cross_entropy_with_logits(
        labels=y_ph_onehot, logits=y_eval)
    net_hooks['eval_output']['loss_eval'] = tf.reduce_mean(loss_eval)
    net_hooks['eval_output']['y_eval'] = y_eval
    net_hooks['placeholder']['Ws_ph'] = Ws_ph
    net_hooks['placeholder']['bs_ph'] = bs_ph
  return net_hooks


# + [markdown] colab_type="text" id="xebqAIwmq9ia"
# # Build model graph

# + 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"
tf.reset_default_graph()

X_ph = tf.placeholder(tf.float32, [None, Hn, Wn, Cn])
y_ph = tf.placeholder(tf.int32, [None])
lr_ph = tf.placeholder(tf.float32)
dropout_rate_ph = tf.placeholder(tf.float32, [])
L2_reg_constant_ph = tf.placeholder(tf.float32, [])

architecture = 'CNN'

if architecture == 'CNN':
  with tf.variable_scope('to_get_shape'):
    # Medium CNN
    filter_sizes = (3, 3, 3, 3)
    pools = (2, 2, 2, 2)
    channels = (64, 128, 256, 256)

    dummy_cnn_hooks = multilayer_CNN(
        X_ph,
        y_ph,
        dropout_rate_ph_in=dropout_rate_ph,
        filter_sizes=filter_sizes,
        pools=pools,
        channels=channels,
        classes=classes)
  # Number of parameters.
  params = dummy_cnn_hooks['weights']['Ws'] + dummy_cnn_hooks['weights']['bs']
  flat_params = tf.concat([tf.reshape(v, [-1]) for v in params], axis=0)
  number_of_params = flat_params.get_shape().as_list()[0]
  print('Number of free parameters=' + str(number_of_params))

  cnn_hooks = multilayer_CNN(
      X_ph,
      y_ph,
      dropout_rate_ph_in=dropout_rate_ph,
      filter_sizes=filter_sizes,
      pools=pools,
      channels=channels,
      classes=classes)
  y = cnn_hooks['train_output']['y']
  loss = cnn_hooks['train_hook']['loss']
  y_eval = cnn_hooks['eval_output']['y_eval']
  pred_eval = cnn_hooks['eval_output']['pred_eval']
  loss_eval = cnn_hooks['eval_output']['loss_eval']
  Ws_ph = cnn_hooks['placeholder']['Ws_ph']
  bs_ph = cnn_hooks['placeholder']['bs_ph']
  Ws = cnn_hooks['weights']['Ws']
  bs = cnn_hooks['weights']['bs']
  L2_reg_Ws = cnn_hooks['train_hook']['L2_reg_Ws']
  L2_reg_bs = cnn_hooks['train_hook']['L2_reg_bs']

loss_to_be_optimized = loss + L2_reg_constant_ph * (L2_reg_Ws + L2_reg_bs)

train_step = tf.train.AdamOptimizer(lr_ph).minimize(loss_to_be_optimized)

correct_prediction = tf.equal(tf.cast(tf.argmax(y, 1), dtype=tf.int32), y_ph)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

correct_prediction_eval = tf.equal(
    tf.cast(tf.argmax(y_eval, 1), dtype=tf.int32), y_ph)
accuracy_eval = tf.reduce_mean(tf.cast(correct_prediction_eval, tf.float32))


# + [markdown] colab_type="text" id="8PHJJ59SuMAo"
# ## Helper Functions

# + colab={} colab_type="code" id="eL52WPi6uPDD"
# Helper functions to flatten weights to a vector, and reform a flattened vector
# to weights.
def flatten(Ws1, bs1):
  lists_now = []
  for W_now in Ws1:
    lists_now.append(W_now.reshape([-1]))
  for b_now in bs1:
    lists_now.append(b_now.reshape([-1]))
  return np.concatenate(lists_now, axis=0)


def get_flat_name(Ws_tf, bs_tf):
  names = []
  for W in Ws_tf:
    names = names + [W.name] * np.prod(W.get_shape().as_list())
  for b in bs_tf:
    names = names + [b.name] * np.prod(b.get_shape().as_list())
  return names


def reform(flat1):
  sofar = 0
  Ws_out_now = []
  bs_out_now = []
  for W in Ws:
    shape_now = W.get_shape().as_list()
    size_now = np.prod(shape_now)
    elements = flat1[sofar:sofar + size_now]
    sofar = sofar + size_now
    Ws_out_now.append(np.array(elements).reshape(shape_now))
  for b in bs:
    shape_now = b.get_shape().as_list()
    size_now = np.prod(shape_now)
    elements = flat1[sofar:sofar + size_now]
    sofar = sofar + size_now
    bs_out_now.append(np.array(elements).reshape(shape_now))
  return Ws_out_now, bs_out_now


# Brier score for evaluating uncertainty performance.
def brier_scores(labels, probs=None, logits=None):
  """Compute elementwise Brier score.

  Args:
    labels: Tensor of integer labels shape [N1, N2, ...]
    probs: Tensor of categorical probabilities of shape [N1, N2, ..., M].
    logits: If `probs` is None, class probabilities are computed as a softmax
      over these logits, otherwise, this argument is ignored.

  Returns:
    Tensor of shape [N1, N2, ...] consisting of Brier score contribution from
    each element. The full-dataset Brier score is an average of these values.
  """
  assert (probs is None) != (
      logits is None), "Exactly one of probs and logits should be None."
  if probs is None:
    probs = softmax(logits, axis=-1)
  nlabels = probs.shape[-1]
  flat_probs = probs.reshape([-1, nlabels])
  flat_labels = labels.reshape([len(flat_probs)])

  plabel = flat_probs[np.arange(len(flat_labels)), flat_labels]
  out = np.square(flat_probs).sum(axis=-1) - 2 * plabel
  return out.reshape(labels.shape) + 1


# + [markdown] colab_type="text" id="W2OrJv2tvdLv"
# # The root training loop (Independent solutions)
# This section could be time-consuming depending on the size of 'points_to_collect' .

# + colab={} colab_type="code" id="T6T6KJVWvNVP"
# Choosing a subset of the train data for faster eval.
N_train_subset = N_val
train_chosen_indices = np.random.choice(
    range(N_train), N_train_subset, replace=False)
X_train_subset = X_train[train_chosen_indices]
y_train_subset = y_train[train_chosen_indices]

# + colab={} colab_type="code" id="Ls0ziO7ivQo4"
# Number of independent solutions to collect.
points_to_collect = 5
# Number of ensemble models to use.
max_ens_size = 4
# Number of full training trajectory to collect for analysis such as T-SNE.
num_trajectory_record = 3
# Collect checkpoints along trajectory for subspace sampling or SWA.
collect_solution_after_epoch = 30

# Training hyperparams for each of the runs.
epochs = 40
batch_size = 128
dropout = 0.1
L2_constant = 0.0

learning_rate = 1.6e-3
# Defining the LR schedule
lr_halving_epoch = 10
lr_halving_maxcount = 1000
lr_halving_power = 2.0

# + colab={} colab_type="code" id="Fw6WhPz-slnf"
print("train_step", train_step)

Ws_many = []
bs_many = []

losses_many = []
accs_many = []

# Collecting last epochs from each trajectory.
Ws_by_epochs_many = [[] for _ in range(points_to_collect)]
bs_by_epochs_many = [[] for _ in range(points_to_collect)]

# Collecting whole trajectory.
Ws_trajectory = [[] for _ in range(num_trajectory_record)]
bs_trajectory = [[] for _ in range(num_trajectory_record)]

global_init_op = tf.global_variables_initializer()
with tf.Session() as sess:
  for point_id in range(points_to_collect):
    print("Optimization " + str(point_id) + " with starting lr=" +
          str(learning_rate) + " dropout rate=" + str(dropout) + " batch=" +
          str(batch_size) + " L2const=" + str(L2_constant))
    sess.run(global_init_op)
    for e in range(epochs):
      iterations = int(np.floor(float(N_train) / float(batch_size)))

      losses_train_list = []
      accs_train_list = []

      halvings_count_now = np.floor(float(e) / lr_halving_epoch)
      halvings_to_be_used = np.min([halvings_count_now, lr_halving_maxcount])
      learning_rate_now = learning_rate / (
          (lr_halving_power)**halvings_to_be_used)

      for it in range(iterations):
        indices = np.random.choice(range(N_train), batch_size)
        X_batch = X_train[indices]
        y_batch = y_train[indices]

        feed_dict = {
            X_ph: X_batch,
            y_ph: y_batch,
            lr_ph: learning_rate_now,
            dropout_rate_ph: dropout,
            L2_reg_constant_ph: L2_constant,
        }

        variables = [train_step, loss, accuracy]

        _, loss_out, acc_out = sess.run(variables, feed_dict=feed_dict)

        losses_train_list.append(loss_out)
        accs_train_list.append(acc_out)

      # Evaluating current epoch.
      feed_dict = {
          X_ph: X_val,
          y_ph: y_val,
          dropout_rate_ph: 0.0,
          L2_reg_constant_ph: 0.0,
      }

      loss_val_out, acc_val_out, Ws_opt_out_now, bs_opt_out_now = sess.run(
          [loss, accuracy, Ws, bs], feed_dict=feed_dict)

      feed_dict[X_ph] = X_train_subset
      feed_dict[y_ph] = y_train_subset

      loss_train_out, acc_train_out = sess.run([loss, accuracy],
                                               feed_dict=feed_dict)

      feed_dict[X_ph] = X_test
      feed_dict[y_ph] = y_test

      loss_test_out, acc_test_out = sess.run([loss, accuracy],
                                             feed_dict=feed_dict)

      print("e=" + str(e) + " train loss=" + f"{loss_train_out:.4f}" +
            " train acc=" + f"{acc_train_out:.4f}" + " val loss=" +
            f"{loss_val_out:.4f}" + " val acc=" + f"{acc_val_out:.4f}" +
            " test loss=" + f"{loss_test_out:.4f}" + " test acc=" +
            f"{acc_test_out:.4f}")
      if e >= collect_solution_after_epoch:
        Ws_by_epochs_many[point_id].append(Ws_opt_out_now)
        bs_by_epochs_many[point_id].append(bs_opt_out_now)
      if point_id < num_trajectory_record:
        Ws_trajectory[point_id].append(Ws_opt_out_now)
        bs_trajectory[point_id].append(bs_opt_out_now)

    # Saving model weights of the last checkpoint.
    Ws_opt_out_now, bs_opt_out_now = sess.run([Ws, bs])

    Ws_many.append(Ws_opt_out_now)
    bs_many.append(bs_opt_out_now)
    losses_many.append(loss_val_out)
    accs_many.append(acc_val_out)

# + [markdown] colab_type="text" id="KapGgJupd9SZ"
# ###Collecting Validation Predictions

# + colab={} colab_type="code" id="8SJntXp5U2Ir"
trajectory_preds_test = np.zeros(
    (num_trajectory_record, epochs, N_val, classes))
independent_preds_test = np.zeros((points_to_collect, N_val, classes))

with tf.Session() as sess:
  for id_now in range(points_to_collect):
    for j, W in enumerate(Ws_many[id_now]):
      feed_dict[Ws_ph[j]] = W
    for j, b in enumerate(bs_many[id_now]):
      feed_dict[bs_ph[j]] = b

    feed_dict[dropout_rate_ph] = 0.0

    feed_dict[X_ph] = X_val
    feed_dict[y_ph] = y_val
    pred_eval_out = sess.run(pred_eval, feed_dict=feed_dict)
    independent_preds_test[id_now] = pred_eval_out

  for id_now in range(num_trajectory_record):
    for e in range(epochs):
      for j, W in enumerate(Ws_trajectory[id_now][e]):
        feed_dict[Ws_ph[j]] = W
      for j, b in enumerate(bs_trajectory[id_now][e]):
        feed_dict[bs_ph[j]] = b
      pred_eval_out = sess.run(pred_eval, feed_dict=feed_dict)
      trajectory_preds_test[id_now][e] = pred_eval_out


# + [markdown] colab_type="text" id="MSyiWg_W9b4F"
# #Cosine and Fraction of Disagreement

# + colab={} colab_type="code" id="wVnTmgAD9l9Y"
def cos_between(v1, v2):
  """ Returns the angle in radians between vectors 'v1' and 'v2'"""
  v1_u = v1 / np.linalg.norm(v1)
  v2_u = v2 / np.linalg.norm(v2)
  return np.dot(v1_u, v2_u)


# + [markdown] colab_type="text" id="FV_a7LZ6BROc"
# ### Within Trajectory

# + colab={} colab_type="code" id="HLjL_tcXGAZp"
trajectory_id = 0
flat_p_list = []
num_epochs = len(Ws_trajectory[trajectory_id])
cos_matrix = np.zeros((num_epochs, num_epochs))

for e in range(num_epochs):
  flat_p_list.append(
      flatten(Ws_trajectory[trajectory_id][e], bs_trajectory[trajectory_id][e]))
for i in range(num_epochs):
  for j in range(i, num_epochs):
    cos_matrix[i][j] = cos_between(flat_p_list[i], flat_p_list[j])
    cos_matrix[j][i] = cos_matrix[i][j]

plt.imshow(cos_matrix, interpolation="nearest", cmap="bwr", origin="lower")
plt.colorbar()
plt.grid("off")

title = "Cosine Along Train Trajectory"
plt.title(title)

plt.xlabel("Checkpoint id")
plt.ylabel("Checkpoint id")
plt.show()

# + colab={} colab_type="code" id="gprmfsa2n8a_"
preds_now = trajectory_preds_test[trajectory_id]
targets_now = y_val

classes_predicted = np.argmax(preds_now, axis=-1)
fractional_differences = np.mean(
    classes_predicted.reshape([1, epochs, len(targets_now)]) !=
    classes_predicted.reshape([epochs, 1, len(targets_now)]),
    axis=-1)

plt.imshow(
    fractional_differences, interpolation="nearest", cmap="bwr", origin="lower")
plt.colorbar()
plt.grid("off")

title = "Disagreement Fraction Along Train Trajectory"
plt.title(title)

plt.xlabel("Checkpoint id")
plt.ylabel("Checkpoint id")
plt.show()

# + [markdown] colab_type="text" id="-FEPT3_OBT-Y"
# ###Between Independent Runs

# + colab={} colab_type="code" id="Wta3MB2tH4jB"
flat_p_list = []
cos_matrix = np.zeros((points_to_collect, points_to_collect))
for i in range(points_to_collect):
  flat_p_list.append(flatten(Ws_many[i], bs_many[i]))

for i in range(points_to_collect):
  for j in range(i, points_to_collect):
    cos_matrix[i][j] = cos_between(flat_p_list[i], flat_p_list[j])
    cos_matrix[j][i] = cos_matrix[i][j]

plt.imshow(cos_matrix, cmap="bwr", origin="lower")
plt.colorbar()
plt.grid("off")

title = "Cosine Between Independent Solutions"
plt.title(title)

plt.xlabel("Independent Solution")
plt.ylabel("Independent Solution")
plt.show()

# + colab={} colab_type="code" id="Cglsoi44tCnL"
preds_now = independent_preds_test
targets_now = y_val

classes_predicted = np.argmax(preds_now, axis=-1)
fractional_differences = np.mean(
    classes_predicted.reshape(
        [1, points_to_collect, len(targets_now)]) != classes_predicted.reshape(
            [points_to_collect, 1, len(targets_now)]),
    axis=-1)

plt.imshow(
    fractional_differences, interpolation="nearest", cmap="bwr", origin="lower")
plt.colorbar()
plt.grid("off")

title = "Disagreement Fraction Btw Independent Solutions"
plt.title(title)

plt.xlabel("Independent Solution")
plt.ylabel("Independent Solution")
plt.show()

# + [markdown] colab_type="text" id="8hgROgxil29I"
# #T-SNE Visualization

# + colab={} colab_type="code" id="W6aVhQ5Yl7A6"
from sklearn.manifold import TSNE

reshaped_prediction = trajectory_preds_test.reshape([-1, N_val * classes])

prediction_embed = TSNE(n_components=2).fit_transform(reshaped_prediction)
traj_embed = prediction_embed.reshape([num_trajectory_record, epochs, 2])

colors_list = ["r", "b", "g"]
labels_list = ["traj_{}".format(i) for i in range(num_trajectory_record)]
for i in range(num_trajectory_record):
  plt.plot(
      traj_embed[i, :, 0],
      traj_embed[i, :, 1],
      color=colors_list[i],
      alpha=0.8,
      linestyle="",
      marker="o",
      label=labels_list[i])
  plt.plot(
      traj_embed[i, :, 0],
      traj_embed[i, :, 1],
      color=colors_list[i],
      alpha=0.3,
      linestyle="-",
      marker="")
  plt.plot(
      traj_embed[i, 0, 0],
      traj_embed[i, 0, 1],
      color=colors_list[i],
      alpha=1.0,
      linestyle="",
      marker="*",
      markersize=20)
plt.legend(loc=1)

# + [markdown] colab_type="text" id="ewL6qSnHlNEy"
# #Effects of Ensemble + Subspace Sampling

# + [markdown] colab_type="text" id="FN3eJMBc55hG"
# ## Gaussian Sampling
# Sample from model weight space according to a Gaussian distribution formed by last epoch checkpoints along a training trajectory.

# + colab={} colab_type="code" id="d5TWmKrc53kK"
from sklearn.decomposition import PCA


def get_gaussian_sample(var_mean, var_std, scale=1.0):
  var_sample = np.random.normal(var_mean, scale * var_std)
  return var_sample


def get_pca_gaussian_flat_sampling(pca, means, rank, scale=1.0):
  standard_normals = np.random.normal(loc=0.0, scale=scale, size=(rank))
  shifts = pca.inverse_transform(standard_normals)
  return shifts + means


def get_rand_norm_direction(shape):
  random_direction = np.random.normal(loc=0.0, scale=1.0, size=shape)
  random_direction_normed = random_direction / np.linalg.norm(random_direction)
  return random_direction_normed


# + colab={} colab_type="code" id="mtGq9Qkj7iEG"
# PCA rank.
rank = 5
num_sample = 30

dial_gaussian_whole_Ws = [[] for _ in range(points_to_collect)]
dial_gaussian_whole_bs = [[] for _ in range(points_to_collect)]

pca_gaussian_whole_Ws = [[] for _ in range(points_to_collect)]
pca_gaussian_whole_bs = [[] for _ in range(points_to_collect)]

for mid in range(points_to_collect):
  Ws_traj = Ws_by_epochs_many[mid]
  bs_traj = bs_by_epochs_many[mid]

  vs_list = []
  for i in range(len(Ws_traj)):
    vs_list.append(flatten(Ws_traj[i], bs_traj[i]))

  vs_np = np.stack(vs_list, axis=0)

  means = np.mean(vs_np, axis=0)
  stds = np.std(vs_np, axis=0)
  vs_np_centered = vs_np - means.reshape([1, -1])

  pca = PCA(n_components=rank)
  pca.fit(vs_np_centered)
  for i in range(num_sample):
    v_sample = get_gaussian_sample(means, stds, scale=1.0)
    w_sample, b_sample = reform(v_sample)
    dial_gaussian_whole_Ws[mid].append(w_sample)
    dial_gaussian_whole_bs[mid].append(b_sample)

    v_sample = get_pca_gaussian_flat_sampling(pca, means, rank, scale=1.0)
    w_sample, b_sample = reform(v_sample)

    pca_gaussian_whole_Ws[mid].append(w_sample)
    pca_gaussian_whole_bs[mid].append(b_sample)

# + [markdown] colab_type="text" id="U_kvSzLKzBSW"
# ## Random Sampling
# Randomly sample from model weight space around the final checkpoint.

# + colab={} colab_type="code" id="PGILtMX7zUHp"
# Random samples need to meet this accuracy threshold to be included.
acc_threshold = 0.70

rand_Ws = [[] for _ in range(points_to_collect)]
rand_bs = [[] for _ in range(points_to_collect)]

with tf.Session() as sess:
  feed_dict = {
    X_ph: X_val,
    y_ph: y_val,
    lr_ph: 0.0,
    dropout_rate_ph: 0.0,
    L2_reg_constant_ph: 0.0,
    }
  for mid in range(points_to_collect):
    for i in range(num_sample):
      vs = flatten(Ws_many[mid], bs_many[mid])  
      for trial in range(5):
        scale = 10*np.random.uniform()   
        v_sample = vs + scale * get_rand_norm_direction(vs.shape)
        w_sample, b_sample = reform(v_sample)

        for j,W in enumerate(w_sample):
          feed_dict[Ws_ph[j]] = w_sample[j]
        for j,b in enumerate(b_sample):
          feed_dict[bs_ph[j]] = b_sample[j]  
      
        val_acc = sess.run(accuracy_eval,feed_dict = feed_dict)
        if val_acc >= acc_threshold:
          rand_Ws[mid].append(w_sample)
          rand_bs[mid].append(b_sample)
          print('Obtaining 1 rand sample at scale {} with validation acc {} at {}th try'.format(scale, val_acc,trial))
          break
        if trial ==4:
          print('No luck -------------------')

# + colab={} colab_type="code" id="NJ5_RCriP-p5"
# PCA low-rank approximation of the random samplings
pca_gaussian_rand_Ws = [[] for _ in range(points_to_collect)]
pca_gaussian_rand_bs = [[] for _ in range(points_to_collect)]

for mid in range(points_to_collect):
  Ws_traj = rand_Ws[mid]
  bs_traj = rand_bs[mid]

  vs_list = []
  for i in range(len(Ws_traj)):
    vs_list.append(flatten(Ws_traj[i], bs_traj[i]))

  vs_np = np.stack(vs_list, axis=0)

  means = np.mean(vs_np, axis=0)
  stds = np.std(vs_np, axis=0)
  vs_np_centered = vs_np - means.reshape([1, -1])

  pca = PCA(n_components=rank)
  pca.fit(vs_np_centered)
  for i in range(num_sample):
    v_sample = get_pca_gaussian_flat_sampling(pca, means, rank, scale=1.0)
    w_sample, b_sample = reform(v_sample)

    pca_gaussian_rand_Ws[mid].append(w_sample)
    pca_gaussian_rand_bs[mid].append(b_sample)


# + [markdown] colab_type="text" id="nEQ96SeTdnkZ"
# ## Collecting predictions from Original and Subspace

# + colab={} colab_type="code" id="gKnPAilsoSNB"
# Average a list of weights.
def average_var(w_list):
  avg = [[] for _ in w_list[0]]
  for w_now in w_list:
    for i, w in enumerate(w_now):
      avg[i].append(w)

  for i, v in enumerate(avg):
    avg[i] = np.mean(np.stack(v, axis=0), axis=0)

  return avg


# Given a list of model weights, feed_dict and a session, returns the model
# predictions as a list.
def get_pred_list(Ws_list, bs_list, feed_dict, sess):
  pred_list = []
  for id in range(len(Ws_list)):
    Ws_now = Ws_list[id]
    bs_now = bs_list[id]
    for j, W in enumerate(Ws):
      feed_dict[Ws_ph[j]] = Ws_now[j]
    for j, b in enumerate(bs):
      feed_dict[bs_ph[j]] = bs_now[j]

    pred_eval_out = sess.run(pred_eval, feed_dict=feed_dict)
    pred_list.append(pred_eval_out)
  return pred_list


# Consider a list of subspaces, each has a list of sampled weights. This
# function computes model predictions, ensembles the predictions within each
# subspace, and returns the list of ensembled predictions.
def get_subspace_pred_list(Ws_subspace_list, bs_subspace_list, feed_dict, sess):
  subspace_pred = []
  num_subspace = len(Ws_subspace_list)
  for mid in range(num_subspace):
    pred_list_now = get_pred_list(Ws_subspace_list[mid], bs_subspace_list[mid],
                                  feed_dict, sess)
    subspace_pred.append(np.mean(np.stack(pred_list_now, axis=0), axis=0))
  return subspace_pred


# + colab={} colab_type="code" id="xeAuHy2xQ39g"
# Returns a list of all possible k-subset from [1, ..., n]
# Don't scale well for large n. Use with caution.
def choose_k_from_n(n, k):
  if k>n or k <1 or n < 1:
    return []
  if k == n :
    return [list(range(1, n+1))]
  if k == 1:
    return [[i] for i in range(1, n+1)]
  a = choose_k_from_n(n-1,k)
  b = choose_k_from_n(n-1, k-1)
  b_new = []
  for g in b:
    b_new.append(g+[n])
  return a+b_new


def get_acc_brier(y, pred, is_logits=False):
  acc = np.mean(np.argmax(pred,axis=-1)==y)
  if is_logits:
    brier = brier_scores(y,logits=pred)
  else:
    brier = brier_scores(y,probs=pred)
  return acc, np.mean(brier)


# Given a list of model predictions, compute the accuracy and brier score for 
# each individual model as well as the ensemble of them.
def get_all_models_metrics(pred_list,y_test,max_ens_size=5, ens_size_list=None):
  acc_list = []
  acc_list_ensemble = []
  b_list = []
  b_list_ensemble = []
  num_models = len(pred_list)
  for i in range(num_models):
    acc,brier = get_acc_brier(y_test, pred_list[i])
    acc_list.append(acc)
    b_list.append(brier)

  max_ens_size = np.min([max_ens_size, num_models])
  if ens_size_list is None:
    ens_size_list = range(1, max_ens_size+1)
  for ens_size in ens_size_list:
    # Pick all possible subset with size of ens_size from available models.
    # Compute ensemble for each such subset.
    ens_index_list = choose_k_from_n(num_models,ens_size)
    ens_acc = []
    ens_brier=[]
    for ens_ind in ens_index_list:
      ens_pred_list = []
      for ind in ens_ind:
        ens_pred_list.append(pred_list[ind-1])
      ens_np = np.mean(np.stack(ens_pred_list,axis=0), axis=0)
      acc,brier = get_acc_brier(y_test, ens_np)
      ens_acc.append(acc)
      ens_brier.append(brier)
    acc_list_ensemble.append(np.mean(ens_acc))
    b_list_ensemble.append(np.mean(ens_brier))
  metrics = {'accuracy': {},
             'brier':{},}
  metrics['accuracy']['individual'] = acc_list
  metrics['accuracy']['ensemble'] = acc_list_ensemble
  metrics['brier']['individual'] = b_list
  metrics['brier']['ensemble'] = b_list_ensemble
  return metrics



# + colab={} colab_type="code" id="RR2UKvMcdpJm"
# Get predictions on test data.
orig_pred = []
wa_pred = []

# Compute averaged weights.
wa_Ws = [[] for _ in range(points_to_collect)]
wa_bs = [[] for _ in range(points_to_collect)]
for i in range(points_to_collect):
  wa_Ws[i] = average_var(Ws_by_epochs_many[i])
  wa_bs[i] = average_var(bs_by_epochs_many[i])

with tf.Session() as sess:
  feed_dict[X_ph] = X_test
  feed_dict[y_ph] = y_test
  orig_pred = get_pred_list(Ws_many, bs_many, feed_dict, sess)
  wa_pred = get_pred_list(wa_Ws, wa_bs, feed_dict, sess)
  diag_gaus_pred = get_subspace_pred_list(dial_gaussian_whole_Ws,
                                          dial_gaussian_whole_bs, feed_dict,
                                          sess)
  pca_gaus_pred = get_subspace_pred_list(pca_gaussian_whole_Ws,
                                         pca_gaussian_whole_bs, feed_dict, sess)
  pca_rand_pred = get_subspace_pred_list(pca_gaussian_rand_Ws,
                                         pca_gaussian_rand_bs, feed_dict, sess)
  rand_pred = get_subspace_pred_list(rand_Ws, rand_bs, feed_dict, sess)

# + colab={} colab_type="code" id="sCSsNdA06lm1"
max_ens_size = points_to_collect - 1

orig_metrics = get_all_models_metrics(
    orig_pred, y_test, max_ens_size=max_ens_size)
wa_metrics = get_all_models_metrics(wa_pred, y_test, max_ens_size=max_ens_size)

diag_metrics = get_all_models_metrics(
    diag_gaus_pred, y_test, max_ens_size=max_ens_size)
pca_metrics = get_all_models_metrics(
    pca_gaus_pred, y_test, max_ens_size=max_ens_size)

rand_metrics = get_all_models_metrics(
    rand_pred, y_test, max_ens_size=max_ens_size)
pca_rand_metrics = get_all_models_metrics(
    pca_rand_pred, y_test, max_ens_size=max_ens_size)

# + [markdown] colab_type="text" id="yQtoJGx817vV"
# ##Plot Metrics for Ensembles + Subspace

# + colab={} colab_type="code" id="utAi4c29CvgO"
title = "Ensemble ACC test"

plt.xlabel("Ensemble size")
plt.ylabel("Test Acc")

ensemble_sizes = np.asarray(range(max_ens_size)) + 1
plt.plot(
    ensemble_sizes,
    orig_metrics["accuracy"]["ensemble"],
    marker="s",
    label="probs ensembling",
    color="navy")

plt.plot(
    ensemble_sizes,
    [np.mean(orig_metrics["accuracy"]["individual"])] * len(ensemble_sizes),
    label="original",
    color="blue")

plt.plot(
    ensemble_sizes,
    [np.mean(diag_metrics["accuracy"]["individual"])] * len(ensemble_sizes),
    label="Diag",
    color="pink")

plt.plot(
    ensemble_sizes,
    [np.mean(pca_metrics["accuracy"]["individual"])] * len(ensemble_sizes),
    label="PCA",
    color="green")

plt.plot(
    ensemble_sizes,
    diag_metrics["accuracy"]["ensemble"],
    marker="s",
    label="Diag+Ensemble",
    color="red")
plt.plot(
    ensemble_sizes,
    wa_metrics["accuracy"]["ensemble"],
    marker="s",
    label="WA+ensembling",
    color="grey")

plt.plot(
    ensemble_sizes,
    pca_metrics["accuracy"]["ensemble"],
    marker="s",
    label="PCA+Ensemble",
    color="green")
plt.plot(
    ensemble_sizes,
    rand_metrics["accuracy"]["ensemble"],
    marker="s",
    label="Rand+Ensemble",
    color="yellow")
plt.plot(
    ensemble_sizes,
    pca_rand_metrics["accuracy"]["ensemble"],
    marker="s",
    label="PCA Rand+Ensemble",
    color="m")

plt.legend()
plt.xlim(1, max_ens_size)
plt.show()

# + colab={} colab_type="code" id="8vLQoHba6eXt"
title = "Ensemble Brier test"

plt.xlabel("Ensemble size")
plt.ylabel("Test Brier")

ensemble_sizes = np.asarray(range(max_ens_size)) + 1
plt.plot(
    ensemble_sizes,
    orig_metrics["brier"]["ensemble"],
    marker="s",
    label="probs ensembling",
    color="navy")

plt.plot(
    ensemble_sizes,
    [np.mean(orig_metrics["brier"]["individual"])] * len(ensemble_sizes),
    label="original",
    color="blue")

plt.plot(
    ensemble_sizes,
    [np.mean(diag_metrics["brier"]["individual"])] * len(ensemble_sizes),
    label="Diag",
    color="pink")

plt.plot(
    ensemble_sizes,
    [np.mean(pca_metrics["brier"]["individual"])] * len(ensemble_sizes),
    label="PCA",
    color="green")

plt.plot(
    ensemble_sizes,
    diag_metrics["brier"]["ensemble"],
    marker="s",
    label="Diag Ensemble",
    color="red")

plt.plot(
    ensemble_sizes,
    pca_metrics["brier"]["ensemble"],
    marker="s",
    label="PCA Ensemble",
    color="green")

plt.plot(
    ensemble_sizes,
    wa_metrics["brier"]["ensemble"],
    marker="s",
    label="WA ensembling",
    color="grey")

plt.plot(
    ensemble_sizes,
    rand_metrics["brier"]["ensemble"],
    marker="s",
    label="Rand Ensemble",
    color="yellow")
plt.plot(
    ensemble_sizes,
    pca_rand_metrics["brier"]["ensemble"],
    marker="s",
    label="PCA Rand Ensemble",
    color="m")

plt.xlim(1, max_ens_size)
plt.legend()
plt.show()

# + [markdown] colab_type="text" id="bGelr5_KBL-O"
# ## Evaluating on Corrupted CIFAR10
# See [paper](https://arxiv.org/pdf/1906.02530.pdf) for data description.
#
# 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].

# + colab={} colab_type="code" id="K2jeOP8AERPn"
ALL_CORRUPTIONS = [
    'gaussian_noise',
    'shot_noise',
    # 'impulse_noise',
    # 'defocus_blur',
    # 'frosted_glass_blur',
    # 'motion_blur',
    # 'zoom_blur',
    # 'snow',
    # 'frost',
    # 'fog',
    # 'brightness',
    # 'contrast',
    # 'elastic',
    # 'pixelate',
    # 'jpeg_compression',
    # 'gaussian_blur',
    # 'saturate',
    # 'spatter',
    # 'speckle_noise',
]

num_models_to_use = 5
ens_size_list = [2, 4]

model_names = ['original', 'wa', 'diag', 'pca', 'rand', 'pca_rand']
model_names_plus_ens = []
for i in ens_size_list:
  model_names_plus_ens = model_names_plus_ens + [
      name + '_ens_{}'.format(i) for name in model_names
  ]

all_model_names = model_names + model_names_plus_ens

model_to_acc_every_level = {name: [] for name in all_model_names}
model_to_brier_every_level = {name: [] for name in all_model_names}

intensity_range = range(6)
for level in intensity_range:
  print('========= Level {} ======='.format(level))
  model_to_acc_this_level = {name: [] for name in all_model_names}
  model_to_brier_this_level = {name: [] for name in all_model_names}
  if level == 0:
    corruptions = ['no_corruption']
  else:
    corruptions = ALL_CORRUPTIONS
  for corruption_type in corruptions:
    if corruption_type is 'no_corruption':
      ds_test = tfds.load(name='cifar10', split=['test'], batch_size=-1)
    else:
      # Load corrupted data.
      corruption_config_name = corruption_type + '_{}'.format(level)
      ds_test = tfds.load(
          name='cifar10_corrupted',
          split=['test'],
          builder_kwargs={'config': corruption_config_name},
          batch_size=-1)
    numpy_ds = tfds.as_numpy(ds_test)
    x_test, y_test = numpy_ds[0]['image'], numpy_ds[0]['label']
    x_test = x_test.reshape([-1, Hn, Wn, Cn])
    y_test = y_test.reshape([-1])
    N_test = len(y_test)
    # Run inference
    with tf.Session() as sess:
      feed_dict[X_ph] = x_test
      feed_dict[y_ph] = y_test
      #Get predictions
      for name in model_names:
        if name is 'original':
          pred_list = get_pred_list(Ws_many[0:num_models_to_use],
                                    bs_many[0:num_models_to_use], feed_dict,
                                    sess)
        elif name is 'wa':
          # Weight Averaged
          pred_list = get_pred_list(wa_Ws[0:num_models_to_use],
                                    wa_bs[0:num_models_to_use], feed_dict, sess)
        elif name is 'diag':
          # Subspacesampling
          pred_list = get_subspace_pred_list(
              dial_gaussian_whole_Ws[0:num_models_to_use],
              dial_gaussian_whole_bs[0:num_models_to_use], feed_dict, sess)
        elif name is 'pca':
          pred_list = get_subspace_pred_list(
              pca_gaussian_whole_Ws[0:num_models_to_use],
              pca_gaussian_whole_bs[0:num_models_to_use], feed_dict, sess)
        elif name is 'pca_rand':
          pred_list = get_subspace_pred_list(
              pca_gaussian_rand_Ws[0:num_models_to_use],
              pca_gaussian_rand_bs[0:num_models_to_use], feed_dict, sess)
        elif name is 'rand':
          pred_list = get_subspace_pred_list(rand_Ws[0:num_models_to_use],
                                             rand_bs[0:num_models_to_use],
                                             feed_dict, sess)
        corrupt_metrics = get_all_models_metrics(
            pred_list, y_test, ens_size_list=ens_size_list)
        model_to_acc_this_level[name].append(
            np.mean(corrupt_metrics['accuracy']['individual']))
        model_to_brier_this_level[name].append(
            np.mean(corrupt_metrics['brier']['individual']))
        for i, ens_size in enumerate(ens_size_list):
          model_to_acc_this_level[name + '_ens_{}'.format(ens_size)].append(
              corrupt_metrics['accuracy']['ensemble'][i])
          model_to_brier_this_level[name + '_ens_{}'.format(ens_size)].append(
              corrupt_metrics['brier']['ensemble'][i])

  for name in all_model_names:
    model_to_acc_every_level[name].append(model_to_acc_this_level[name])
    model_to_brier_every_level[name].append(model_to_brier_this_level[name])

# + [markdown] colab_type="text" id="rGS9GPU2mo5j"
# ## Plot the metrics across corruption intensity

# + colab={} colab_type="code" id="b5cAs29LkPYu"
ens_size_list = [2, 4]
model_names = ['original', 'wa', 'diag', 'pca', 'rand', 'pca_rand']
model_names_plus_ens = []
for i in ens_size_list:
  model_names_plus_ens = model_names_plus_ens + [
      name + '_ens_{}'.format(i) for name in model_names
  ]

all_model_names = model_names + model_names_plus_ens

# Model names used by Jensen Shannon plots.
js_model_names = ['inde', 'traj', 'diag', 'pca', 'rand', 'pca_rand']

model_name_to_color = {
    'original': 'indianred',
    'wa': 'dimgray',
    'diag': 'gold',
    'pca': 'blue',
    'rand': 'mediumseagreen',
    'pca_rand': 'fuchsia',
    'inde': 'red',
    'traj': 'grey'
}

for name in model_names:
  for i in ens_size_list:
    model_name_to_color[name + '_ens_{}'.format(i)] = model_name_to_color[name]

model_name_to_label = {
    'original': 'Original',
    'wa': 'Weight Avg',
    'diag': 'Diag Gaus',
    'pca': 'PCA Gaus',
    'rand': 'Random',
    'pca_rand': 'PCA Random',
    'inde': 'Independent',
    'traj': 'Train Trajectory'
}

for name in model_names:
  model_name_to_label[name + '_ens'] = model_name_to_label[name] + '+Ens'
  for i in ens_size_list:
    model_name_to_label[
        name +
        '_ens_{}'.format(i)] = model_name_to_label[name] + '+Ens {}'.format(i)

# Line style, currently can only support 2 ensemble sizes. Beyond that, solid
# line will be used for all.
model_name_to_ls = {}
for name in model_names:
  model_name_to_ls[name] = 'dashed'
  model_name_to_ls[name + '_ens_{}'.format(ens_size_list[0])] = 'dotted'
  for i in range(1, len(ens_size_list)):
    model_name_to_ls[name + '_ens_{}'.format(ens_size_list[i])] = 'solid'


def plot_metric_over_corrupted_data(metric_name, intensity_range,
                                    all_model_names, base_model_names,
                                    color_map, ls_map):
  plt.figure(figsize=(10, 9))

  ylabel = {
      'acc': 'Accuracy',
      'brier': 'Brier',
  }
  plt.xlabel('Corruption Intensity', fontsize=16)
  plt.ylabel(ylabel[metric_name], fontsize=16)
  label_map = {'original': 'Single'}
  for ens_size in ens_size_list:
    label_map['original_ens_{}'.format(ens_size)] = '{}-Ensemble'.format(
        ens_size)
  model_to_mean = {name: [] for name in all_model_names}
  for name in all_model_names:
    if metric_name is 'acc':
      model_to_metric_every_level = model_to_acc_every_level[name]
    elif metric_name is 'brier':
      model_to_metric_every_level = model_to_brier_every_level[name]
    for t in intensity_range:
      model_to_mean[name].append(np.mean(model_to_metric_every_level[t]))

    if name in label_map:
      label = label_map[name]
    else:
      label = None
    plt.plot(
        intensity_range,
        model_to_mean[name],
        marker='s',
        label=label,
        color=color_map[name],
        ls=ls_map[name])

  legend0 = plt.legend(loc=3, fontsize=12)

  patch_list = []
  for name in base_model_names:
    patch_list.append(
        mpatch.Patch(
            color=model_name_to_color[name],
            label=model_name_to_label[name]))
  legend1 = plt.legend(handles=patch_list)
  ax = plt.gca().add_artist(legend0)

  plt.xlim(np.min(intensity_range), np.max(intensity_range))
  plt.tight_layout()


def plot_js_over_corrupted_data(intensity_range, all_model_names, color_map,
                                label_map):
  plt.figure(figsize=(10, 9))
  plt.xlabel('Corruption Intensity', fontsize=16)
  plt.ylabel('<NAME>', fontsize=16)

  model_to_mean = {name: [] for name in all_model_names}

  for name in all_model_names:
    model_to_metric_every_level = model_to_js_every_level[name]
    for t in intensity_range:
      model_to_mean[name].append(np.mean(model_to_metric_every_level[t]))
    plt.plot(
        intensity_range,
        model_to_mean[name],
        marker='s',
        label=label_map[name],
        color=color_map[name])

  plt.xlim(np.min(intensity_range), np.max(intensity_range))
  plt.tight_layout()
  plt.legend()


# + colab={} colab_type="code" id="cszvIYjZkHLi"
intensity_range = range(6)
plot_metric_over_corrupted_data('acc', intensity_range, all_model_names,
                                model_names, model_name_to_color,
                                model_name_to_ls)

# + colab={} colab_type="code" id="aLsVQu1SYzaX"
intensity_range = range(6)
plot_metric_over_corrupted_data('brier', intensity_range, all_model_names,
                                model_names, model_name_to_color,
                                model_name_to_ls)


# + [markdown] colab_type="text" id="ZwOzmfhn3Z4l"
# ##<NAME>

# + colab={} colab_type="code" id="tqiAL-z73cPe"
def get_entropy(p):
  return np.sum(-p * np.log(p + 1e-5), axis=1)


# Compute the Jensen Shannon score given a list of model predictions.
def get_jensen_shannon(pred_list, is_logit=False):
  n = len(pred_list)
  if is_logit:
    p_list = []
    for i in range(n):
      p_list.append(softmax(pred_list[i], axis=1))
  else:
    p_list = pred_list
  p_np = np.stack(p_list, axis=0)
  ensemble_pred = np.mean(p_np, axis=0)
  ensemble_entropy = np.mean(get_entropy(ensemble_pred))

  mean_entropy = np.mean(np.sum(-p_np * np.log(p_np + 1e-5), axis=2))
  return ensemble_entropy - mean_entropy


# + [markdown] colab_type="text" id="mCgvOt_gbRYG"
# ### Eval on SVHN
# JS on OOD data

# + colab={} colab_type="code" id="4WralaRlbVOW"
_, ds_test = tfds.load(
    name="svhn_cropped", split=["train", "test"], batch_size=-1)

numpy_ds = tfds.as_numpy(ds_test)
X_test_raw, y_test_fine_raw = numpy_ds["image"], numpy_ds["label"]

X_test_svhn = X_test_raw
y_test_svhn = y_test_fine_raw

subsample = np.random.choice(len(y_test_svhn), 10000, replace=False)
X_svhn = X_test_svhn[subsample, :, :, :]
y_svhn = y_test_svhn[subsample]

# + colab={} colab_type="code" id="ICYLKUto1mrv"
traj_svhn_js = []
pca_svhn_js = []
diag_svhn_js = []
rand_svhn_js = []
pca_rand_svhn_js = []

with tf.Session() as sess:
  feed_dict[X_ph] = X_svhn
  feed_dict[y_ph] = y_svhn
  feed_dict[lr_ph] = 0.0
  feed_dict[dropout_rate_ph] = 0.0
  indi_svhn_pred = get_pred_list(Ws_many[0:10], bs_many[0:10], feed_dict, sess)
  for traj_id in range(points_to_collect):
    trajectory_svhn_pred = get_pred_list(Ws_by_epochs_many[traj_id],
                                         bs_by_epochs_many[traj_id], feed_dict,
                                         sess)
    traj_svhn_js.append(get_jensen_shannon(trajectory_svhn_pred))

    diag_gaus_whole_svhn_pred = get_pred_list(dial_gaussian_whole_Ws[traj_id],
                                              dial_gaussian_whole_bs[traj_id],
                                              feed_dict, sess)
    diag_svhn_js.append(get_jensen_shannon(diag_gaus_whole_svhn_pred))

    pca_gaus_whole_svhn_pred = get_pred_list(pca_gaussian_whole_Ws[traj_id],
                                             pca_gaussian_whole_bs[traj_id],
                                             feed_dict, sess)
    pca_svhn_js.append(get_jensen_shannon(pca_gaus_whole_svhn_pred))

    pca_rand_svhn_pred = get_pred_list(pca_gaussian_rand_Ws[traj_id],
                                       pca_gaussian_rand_bs[traj_id], feed_dict,
                                       sess)
    pca_rand_svhn_js.append(get_jensen_shannon(pca_rand_svhn_pred))

    rand_svhn_pred = get_pred_list(rand_Ws[traj_id], rand_bs[traj_id],
                                   feed_dict, sess)
    rand_svhn_js.append(get_jensen_shannon(rand_svhn_pred))

# + colab={} colab_type="code" id="-1j-4_KhAZEg"
print("Jensen-Shannon for Independent run is {:.4f}".format(
    get_jensen_shannon(indi_svhn_pred)))
print("Jensen-Shannon for within-trajectory is {:.4f}".format(
    np.mean(traj_svhn_js)))

print("Jensen-Shannon for Diag Gaussian is {:.4f}".format(
    np.mean(diag_svhn_js)))
print("Jensen-Shannon for PCA Gaussian is {:.4f}".format(np.mean(pca_svhn_js)))

print("Jensen-Shannon for Rand is {:.4f}".format(np.mean(rand_svhn_js)))
print("Jensen-Shannon for PCA Rand is {:.4f}".format(np.mean(pca_rand_svhn_js)))

# + [markdown] colab_type="text" id="FLCIw3GwQeaa"
# ###Eval on Cifar-10-C
#
# 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].

# + colab={} colab_type="code" id="UNM4VieiQjAp"
# Choose any training trajectory to study the JS of subspace samplings.
traj_id = 3

model_to_js_every_level = {name: [] for name in js_model_names}

intensity_range = range(6)
for level in intensity_range:
  print("========= Level {} =======".format(level))
  model_to_js_this_level = {name: [] for name in js_model_names}
  if level == 0:
    corruptions = ["no_corruption"]
  else:
    corruptions = ALL_CORRUPTIONS
  for corruption_type in corruptions:
    if corruption_type is "no_corruption":
      ds_test = tfds.load(name="cifar10", split=["test"], batch_size=-1)
    else:
      # Load corrupted data.
      corruption_config_name = corruption_type + "_{}".format(level)
      ds_test = tfds.load(
          name="cifar10_corrupted",
          split=["test"],
          builder_kwargs={"config": corruption_config_name},
          batch_size=-1)
    numpy_ds = tfds.as_numpy(ds_test)
    x_test, y_test = numpy_ds[0]["image"], numpy_ds[0]["label"]
    x_test = x_test.reshape([-1, Hn, Wn, Cn])
    y_test = y_test.reshape([-1])
    N_test = len(y_test)
    # Run inference
    with tf.Session() as sess:
      feed_dict[X_ph] = x_test
      feed_dict[y_ph] = y_test
      #Get predictions
      for name in js_model_names:
        if name is "inde":
          pred_list = get_pred_list(Ws_many, bs_many, feed_dict,
                                    sess)
        elif name is "traj":
          # Weight Averaged
          pred_list = get_pred_list(Ws_by_epochs_many[traj_id],
                                    bs_by_epochs_many[traj_id], feed_dict, sess)
        elif name is "diag":
          pred_list = get_pred_list(dial_gaussian_whole_Ws[traj_id],
                                    dial_gaussian_whole_bs[traj_id], feed_dict,
                                    sess)
        elif name is "pca":
          pred_list = get_pred_list(pca_gaussian_whole_Ws[traj_id],
                                    pca_gaussian_whole_bs[traj_id], feed_dict,
                                    sess)
        elif name is "pca_rand":
          pred_list = get_pred_list(pca_gaussian_rand_Ws[traj_id],
                                    pca_gaussian_rand_bs[traj_id], feed_dict,
                                    sess)
        elif name is "rand":
          pred_list = get_pred_list(rand_Ws[traj_id], rand_bs[traj_id],
                                    feed_dict, sess)

        model_to_js_this_level[name].append(get_jensen_shannon(pred_list))

  for name in js_model_names:
    model_to_js_every_level[name].append(model_to_js_this_level[name])

# + colab={} colab_type="code" id="2mcbo-46d4s2"
plot_js_over_corrupted_data(intensity_range, js_model_names,
                            model_name_to_color, model_name_to_label)

# + [markdown] colab_type="text" id="zCZUtnjKljz-"
# #Diversity vs Accuracy
# This section plots the Diversiy vs Accuracy plane for different subspace sampling methods, compared to independent solutions:
#
#
# *   Random
# *   Gaussian Diag
# *   Gaussian Low rank (PCA)
# *   Dropout
#
# 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

# + colab={} colab_type="code" id="D3eKK8UPz-UB"
base_id = 0
base_v = flatten(Ws_many[base_id], bs_many[base_id])


def get_acc_and_diff(pred_class, base_pred, y_label):
  diff = np.mean(pred_class != base_pred)
  acc = np.mean(pred_class == y_label)
  return acc, diff


# Given a list of weights, compute their accuracy and their difference to the
# base prediction.
def get_acc_and_diff_from_weights(Ws_list, bs_list, feed_dict, base_pred):
  assert len(Ws_list) == len(bs_list)
  with tf.Session() as sess:
    pred_list = get_pred_list(Ws_list, bs_list, feed_dict, sess)

  acc_list = []
  diff_list = []
  for i in range(len(Ws_list)):
    acc, diff = get_acc_and_diff(
        np.argmax(pred_list[i], axis=-1), base_pred, feed_dict[y_ph])
    acc_list.append(acc)
    diff_list.append(diff)
  return acc_list, diff_list


# + [markdown] colab_type="text" id="0oScJ4X09-F9"
# ##Independent Solutions

# + colab={} colab_type="code" id="TkwIdoOJvgz2"
feed_dict = {
    X_ph: X_test,
    y_ph: y_test,
    lr_ph: 0.0,
    dropout_rate_ph: 0.0,
    L2_reg_constant_ph: 0.0,
}

with tf.Session() as sess:
  orig_pred_list = get_pred_list(Ws_many, bs_many, feed_dict, sess)

base_pred = np.argmax(orig_pred_list[base_id], axis=-1)

independent_acc = []
independent_diff = []
for i in range(points_to_collect):
  class_pred = np.argmax(orig_pred_list[i], axis=-1)
  acc, diff = get_acc_and_diff(class_pred, base_pred, y_test)
  independent_acc.append(acc)
  independent_diff.append(diff)

data_to_show = [
    ("independent optima", "red", 300, "*", 1.0, independent_acc,
     independent_diff),
]

# + colab={} colab_type="code" id="cotIEVchK3Vh"
data_to_show = [
    ("independent optima", "red", 300, "*", 1.0, independent_acc,
     independent_diff),
]

# + [markdown] colab_type="text" id="hT1pKG9W-BTZ"
# ##Gaussian Sampling

# + colab={} colab_type="code" id="czoWhgyGAFoC"
num_sample = 50
rank = 5  # for PCA

Ws_traj = Ws_by_epochs_many[base_id]
bs_traj = bs_by_epochs_many[base_id]

dial_gaussian_whole_Ws = []
dial_gaussian_whole_bs = []

pca_gaussian_whole_Ws = []
pca_gaussian_whole_bs = []

vs_list = []
for i in range(len(Ws_traj)):
  vs_list.append(flatten(Ws_traj[i], bs_traj[i]))

vs_np = np.stack(vs_list, axis=0)

means = np.mean(vs_np, axis=0)
stds = np.std(vs_np, axis=0)
vs_np_centered = vs_np - means.reshape([1, -1])

pca = PCA(n_components=rank)
pca.fit(vs_np_centered)
for i in range(num_sample):
  scale = np.random.uniform()
  # One can adjust the constant in front of scale to get a fuller range in
  # diversity-acc plane.
  v_sample = get_gaussian_sample(means, stds, scale=10.0 * scale)
  w_sample, b_sample = reform(v_sample)
  dial_gaussian_whole_Ws.append(w_sample)
  dial_gaussian_whole_bs.append(b_sample)
  # One can adjust the scale value to get a fuller range in diversity-acc plane.
  v_sample = get_pca_gaussian_flat_sampling(pca, means, rank, scale=6.0)
  w_sample, b_sample = reform(v_sample)

  pca_gaussian_whole_Ws.append(w_sample)
  pca_gaussian_whole_bs.append(b_sample)

pca_acc, pca_diff = get_acc_and_diff_from_weights(pca_gaussian_whole_Ws,
                                                  pca_gaussian_whole_bs,
                                                  feed_dict, base_pred)
diag_acc, diag_diff = get_acc_and_diff_from_weights(dial_gaussian_whole_Ws,
                                                    dial_gaussian_whole_bs,
                                                    feed_dict, base_pred)

# + colab={} colab_type="code" id="mfnrhHSa5Xiw"
data_to_show = data_to_show + [
    ("diag gaussian", "purple", 10, "o", 0.6, diag_acc, diag_diff),
    ("pca gaussian", "fuchsia", 10, "o", 0.3, pca_acc, pca_diff),
]

# + [markdown] colab_type="text" id="Jn-g6-GZZ5tc"
# ##Random Sampling

# + colab={} colab_type="code" id="uJpRJaDLaFRa"
rand_Ws = []
rand_bs = []
for i in range(num_sample):
  # One can adjust the constant in front of scale to get a fuller range in
  # diversity-acc plane. 
  scale = 70.0 * np.random.uniform()
  v_sample = base_v + scale * get_rand_norm_direction(base_v.shape)
  w_sample, b_sample = reform(v_sample)
  rand_Ws.append(w_sample)
  rand_bs.append(b_sample)

rand_acc, rand_diff = get_acc_and_diff_from_weights(rand_Ws, rand_bs, feed_dict,
                                                    base_pred)

# + colab={} colab_type="code" id="h1my3OECbzEi"
data_to_show = data_to_show + [
    ("random subspace", "blue", 3, "o", 0.3, rand_acc, rand_diff),
]


# + [markdown] colab_type="text" id="JJWIEJZBZ9Jf"
# ##Dropout Sampling

# + colab={} colab_type="code" id="MBiCleCgb8r1"
def apply_dropout_to_array(anchor_array, dropout_to_use):
  shape_now = anchor_array.shape
  random_mask = np.random.rand(*shape_now)
  mask = (random_mask < (1.0 - dropout_to_use))
  array_dropped = (anchor_array * mask) / (1.0 - dropout_to_use)
  return array_dropped


dropout_Ws = []
dropout_bs = []
for i in range(num_sample):
  # One can adjust the constant to get a fuller range in diversity-acc plane.
  dropout = 0.1 * np.random.uniform()
  v_sample = apply_dropout_to_array(base_v, dropout)
  w_sample, b_sample = reform(v_sample)
  dropout_Ws.append(w_sample)
  dropout_bs.append(b_sample)

dropout_acc, dropout_diff = get_acc_and_diff_from_weights(
    dropout_Ws, dropout_bs, feed_dict, base_pred)

# + colab={} colab_type="code" id="E5qdldWHb9c4"
data_to_show = data_to_show + [
    ("dropout subspace", "orange", 10, "o", 1.0, dropout_acc, dropout_diff),
]


# + [markdown] colab_type="text" id="VwNL8PPM7DrM"
# ##Plotting Figure

# + colab={} colab_type="code" id="AqhcI78L52hM"
# Functions for computing the analytic limit curves (see paper).

def perturbed_reference_analytic(desired_accuracy, reference_accuracy, classes):
  return (reference_accuracy - desired_accuracy) / (
      reference_accuracy + (reference_accuracy - 1.0) / (classes - 1.0))


def random_average_case(desired_accuracy, reference_accuracy, classes):
  part1 = reference_accuracy * (1.0 - desired_accuracy)
  part2 = desired_accuracy * (1.0 - reference_accuracy)
  part3 = (1.0 - reference_accuracy) * (1.0 -
                                        desired_accuracy) * (classes - 2.0) / (
                                            classes - 1.0)
  return part1 + part2 + part3


# + colab={} colab_type="code" id="DjOAuW3F56qu"
plt.figure(figsize=(7, 5))
accs_fit_random = np.linspace(0.1, 0.71, 100)
diffs_fit_random = random_average_case(accs_fit_random, 0.71, 10)
plt.plot(
    accs_fit_random,
    diffs_fit_random / (1.0 - accs_fit_random),
    color="black",
    linestyle="-.",
    label="Upper limit")

# Analytic limits curve.
accs_fit_perturbed = np.linspace(0.1, 0.71, 100)
diffs_fit_perturbed = perturbed_reference_analytic(accs_fit_perturbed, 0.71, 10)
plt.plot(
    accs_fit_perturbed,
    diffs_fit_perturbed / (1.0 - accs_fit_perturbed),
    color="black",
    linestyle="--",
    label="Lower limit")

for (name_now, color_now, size_now, marker_now, alpha_now, accs_now,
     diffs_now) in data_to_show:

  metric_now = np.asarray(diffs_now) / (1.0 - np.asarray(accs_now))

  if name_now != "independent optima":
    plt.scatter(
        accs_now,
        metric_now,
        color=color_now,
        s=size_now,
        marker=marker_now,
        alpha=alpha_now)
    plt.scatter([], [],
                color=color_now,
                label=name_now,
                s=10,
                marker=marker_now,
                alpha=1.0)
  else:
    plt.scatter(
        accs_now,
        metric_now,
        color=color_now,
        s=size_now,
        marker=marker_now,
        alpha=alpha_now,
        label=name_now)

  if name_now == "independent optima":
    base_star_color = "green"
    plt.scatter([accs_now[base_id]], [metric_now[base_id]],
                color=base_star_color,
                label="baseline optimum",
                s=size_now,
                marker=marker_now,
                alpha=alpha_now)

plt.xlabel("Validation accuracy", fontsize=14)
plt.ylabel("Fraction of labels changes / (1.0-accuracy)", fontsize=14)

plt.title("MediumCNN on Cifar10", fontsize=18)
plt.legend(loc=3, fancybox=True, framealpha=0.5, fontsize=16)
plt.ylim([-0.1, 1.2])
plt.xlim([0, 0.78])

plt.show()

StarcoderdataJupyterScriptsDedupFiltered