<jupyter_start><jupyter_text>Machine Learning Seminars - Subspace Inference for Bayesian Deep Learning - Demo*Reviewed by <NAME>, <NAME> and <NAME>***Aims of this notebook**:- Empirically evaluate the effectiveness of the uncertainty estimates produced (this is **not** evaluated by the paper)- Evaluate the performance gains achieved by applying approximate inference techniques in subspaces *vs* in the full space**Research questions**:- In comparison to ensembles and 'full' variational inference (VI); how well does VI in PCA subspace perform in terms of uncertainty estimates (measured by [expected calibration error](https://arxiv.org/pdf/1706.04599.pdf) and out-of-distribution uncertainty estimates).- Using the same comparisons, how does 'performance' of the paper stack up?   Setup<jupyter_code>!rm -rf drbayes
!git clone https://github.com/ymohamedahmed/drbayes.git
import sys
%pip install -e drbayes 
import sys
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

import torch
import torch.utils.data
from torch.nn import functional as F
import torch.nn as nn
from torchvision import datasets, transforms

import seaborn as sns

import subspace_inference
import subspace_inference.utils as utils
from subspace_inference.posteriors import SWAG
from subspace_inference import models, losses, utils
from subspace_inference.models import MLP
from subspace_inference.visualization import plot_predictive
from subspace_inference.posteriors.proj_model import SubspaceModel
from tqdm.notebook import tqdm
from subspace_inference.posteriors.vi_model import VIModel, ELBO
import math

import os

torch.backends.cudnn.benchmark = True
INITIAL_SEED = 1
torch.manual_seed(INITIAL_SEED)
torch.cuda.manual_seed(INITIAL_SEED)
np.random.seed(INITIAL_SEED)

%load_ext autoreload
%autoreload 2
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
train_all = datasets.QMNIST(root="../data",train=True, download=True,
                   transform=transform)
train_val_split = int(0.8*len(train_all))
train_idx, val_idx = range(0, train_val_split), range(train_val_split, len(train_all))

train = torch.utils.data.Subset(train_all, train_idx)
val = torch.utils.data.Subset(train_all, val_idx)
test = datasets.QMNIST('../data', train=False,download=True,
                   transform=transform)
train_loader = torch.utils.data.DataLoader(train,batch_size=512)
val_loader = torch.utils.data.DataLoader(val,batch_size=len(val))<jupyter_output>Downloading https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-images-idx3-ubyte.gz to ../data/QMNIST/raw/qmnist-train-images-idx3-ubyte.gz<jupyter_text>Helper code<jupyter_code>class VanillaMLP(nn.Module):
  def __init__(self, dims):
    super(VanillaMLP,self).__init__()
    self.dims = dims
    layers = [nn.Flatten()] + [lay for (x,y) in zip(dims[:-1],dims[1:]) for lay in [nn.Linear(x,y), nn.ReLU()] ]
    layers.pop()
    self.model = nn.Sequential(*layers)

  def forward(self,x):
    return self.model(x)
def celoss_acc(model, loss_function, loader):
  N = 0
  total_loss = 0
  acc = 0
  for (x,y) in loader:
    x = x.to(device)
    y = y.to(device)
    out = model(x)
    loss = loss_function(model,x,y)[0]
    total_loss += loss.item()
    acc += (out.argmax(dim=1)==y).sum().item()
    N += x.size(0)
  return total_loss/N, acc/N

def convergence(val_losses):
  return len(val_losses) > 10 and sum([val_losses[i] >= val_losses[i-1] for i in range(-1,-10,-1)])>5

def train(model, loss_function, max_epochs, train_loader, val_loader, lr, collect_model=False):
  optimizer = torch.optim.Adam(model.parameters(), lr=lr)
  vlosses = []
  if collect_model:
    swag_model = SWAG(VanillaMLP, subspace_type="pca",
                    subspace_kwargs={"max_rank": 10, "pca_rank": 10},dims=model.dims)
  pbar = tqdm(range(max_epochs))
  for epoch in pbar:
    for x,y in train_loader:
      x = x.to(device)
      y = y.to(device)
      optimizer.zero_grad()
      if collect_model: 
        swag_model.collect_model(model)
      loss = loss_function(model,x,y)[0]
      loss.backward()
      optimizer.step()
    vloss, vacc = celoss_acc(model, loss_function, val_loader)
    vlosses.append(vloss)

    # don't check convergence if simply collecting weights
    if convergence(vlosses) and not(collect_model):
      print("converged")
      break
    pbar.set_description(f'loss: {round(vloss,3)}, acc: {round(vacc,3)}')
    if epoch % 2 == 0:
        print(f"Epoch: {epoch}, loss: {vloss}, acc: {vacc}")

  if collect_model:
    return swag_model.get_space()
def vi_pred(model, number_samples, test_loader, num_classes=10):
  preds = torch.zeros((len(test_loader.dataset), num_classes))
  for n, (x,_) in tqdm(enumerate(test_loader)):
    x = x.to(device)
    for i in range(number_samples):
      preds[n*test_batch_size:(n+1)*test_batch_size] += F.softmax(model(x).cpu(),dim=1)
    del x 
    torch.cuda.empty_cache()
  preds = 1/number_samples *preds
  acc = (1.*torch.argmax(preds, dim=1).eq(test_loader.dataset.targets[:,0])).mean()
  return preds, acc
def number_of_params(dims):
  d = np.array(dims)
  return ((d[:-1]*d[1:])+d[1:]).sum()

def create_full_vi_model():
  init_sigma = 1.
  prior_sigma = 1.
  dims = [28*28,10]
  number_params = number_of_params(dims)
  identity_space = SubspaceModel(torch.zeros(number_params), torch.eye(number_params))
  mod = VIModel(
      subspace=identity_space,
      init_inv_softplus_sigma=math.log(math.exp(init_sigma) - 1.0),
      prior_log_sigma=math.log(prior_sigma),
      base=VanillaMLP,
      dims=dims
  )
  mod.to(device)
  return mod
class Ensemble(nn.Module):
    
  def __init__(self, number_of_networks, dims):
    self.N = number_of_networks
    self.dims = dims
    self.ensemble = []
    for _ in range(number_of_networks):
       # change the seed otherwise all models will start with same init.
       self._new_seeds()
       self.ensemble.append(VanillaMLP(dims))
    self._reset_seeds()

  def _reset_seeds(self):
    torch.manual_seed(INITIAL_SEED)
    torch.cuda.manual_seed(INITIAL_SEED)
    np.random.seed(INITIAL_SEED)
  
  def _new_seeds(self):
    torch.manual_seed(torch.randint(high=256,size=(1,)).item())
    torch.cuda.manual_seed(torch.randint(high=256,size=(1,)).item())
    np.random.seed(torch.randint(high=256,size=(1,)).item())

  def train(self, loss_function, max_epochs, train_loader, val_loader, lr):
    for m in self.ensemble:
      m.to(device)
      self._new_seeds()
      train(m, loss_function, max_epochs, train_loader, val_loader, lr=lr)
    self._reset_seeds()
  
  def _single_predict(self, x):
    output = torch.zeros(self.dims[-1])
    for m in self.ensemble:
      output += m(x)
    return output/self.N
  
  def predict(self, loader):
    preds = torch.zeros(len(loader), self.dims[-1])
    for i, (x,y) in enumerate(loader):
      preds[i] = self._single_predict(x)
    return preds

  def save(self):
    for i in range(self.N):
      torch.save(self.ensemble[i].state_dict(), f'/content/drive/My Drive/sibdl/ensemble-{i}.pt')
    
  def load(self):
    for i in range(self.N):
      self.ensemble[i].load_state_dict(torch.load(f'/content/drive/My Drive/sibdl/ensemble-{i}.pt'))
# run VI on the PCA-derived subspace
def create_vi_model(subspace):
  init_sigma = 1.
  prior_sigma = 5.
  criterion = losses.cross_entropy
  
  vi_model = VIModel(
      subspace=subspace,
      init_inv_softplus_sigma=math.log(math.exp(init_sigma) - 1.0),
      prior_log_sigma=math.log(prior_sigma),
      base=VanillaMLP,
      dims=[28*28,10]
  )
  
  vi_model.to(device)
  return vi_model<jupyter_output><empty_output><jupyter_text>Training the models (you can skip this!)- Full VI - VI in PCA subspace - ensemble of NNs  Mean Field Variational Inference<jupyter_code>full_vi_model = create_full_vi_model()
temperature = 1.
elbo = ELBO(losses.cross_entropy, len(train_loader.dataset), temperature=temperature)
train(full_vi_model, elbo, 500, train_loader, val_loader, lr=0.1)
torch.save(full_vi_model.state_dict(), "/content/drive/My Drive/sibdl/full-vi-pca-model.pt")<jupyter_output><empty_output><jupyter_text>Ensembles<jupyter_code>ensemble = Ensemble(number_of_networks = 10, dims = [28*28,10])
ensemble.train(losses.cross_entropy, 500, train_loader, val_loader, lr=0.01)
ensemble.save()<jupyter_output><empty_output><jupyter_text>PCA subspace<jupyter_code>model = VanillaMLP([28*28,10])
model.to(device)
# train model to convergence

train(model, losses.cross_entropy, 500, train_loader, val_loader, lr=0.001)

# train for a further number of epochs to collect PCA subspace
space = train(model, losses.cross_entropy, 200, train_loader, val_loader, lr=0.001, collect_model=True)
torch.save(model.state_dict(), "/content/drive/My Drive/sibdl/MLP-qmnist.pt")
mean, _, cov_factor = space
subspace = SubspaceModel(mean, cov_factor)
torch.save(subspace.state_dict(), "/content/drive/My Drive/sibdl/pca-subspace-qmnist.pt")
vi_model = create_vi_model(subspace)
temperature = 1.
elbo = ELBO(losses.cross_entropy, len(train_loader.dataset), temperature=temperature)
train(vi_model, elbo, 500, train_loader, val_loader, lr=0.001)
torch.save(vi_model.state_dict(), "/content/drive/My Drive/sibdl/vi-pca-model.pt")<jupyter_output><empty_output><jupyter_text>Running Ensembles, PCA subspace and full VI Let's load the old models<jupyter_code>full_vi_model = create_full_vi_model()
full_vi_model.load_state_dict(torch.load('/content/drive/My Drive/sibdl/full-vi-pca-model.pt'))

ensemble = Ensemble(10, [28*28,10])
ensemble.load()

mean,cov_factor = torch.zeros(7850), torch.zeros(10,7850)
subspace = SubspaceModel(mean, cov_factor)
subspace.load_state_dict(torch.load("/content/drive/My Drive/sibdl/pca-subspace-qmnist.pt"))
vi_model = create_vi_model(subspace)<jupyter_output><empty_output><jupyter_text>Full VI<jupyter_code>test_batch_size = 512
test_loader = torch.utils.data.DataLoader(test,batch_size=test_batch_size)
full_vi_preds,acc = vi_pred(full_vi_model, 100, test_loader)
inds = torch.max(full_vi_preds,dim=1)[0].argsort()[:20]
print(torch.max(full_vi_preds,dim=1)[0][inds])
fig,axs = plt.subplots(figsize=(10,3),nrows=4, ncols=5)
axs = axs.flatten()
for i,im in enumerate(test_loader.dataset.data[inds]):
  ax = axs[i]
  ax.imshow(im,cmap='gray')
  ax.set_axis_off()
  # ax.title.set_text(f'predicted: {full_vi_preds[i].argmax()}')
# plt.title("Full Variational Inference")
plt.savefig('full-vi.png')
plt.show()<jupyter_output><empty_output><jupyter_text>PCA subspace<jupyter_code>test_batch_size = 512
test_loader = torch.utils.data.DataLoader(test,batch_size=test_batch_size)
# vi_preds,acc = vi_pred(vi_model, 100, test_loader)
inds = torch.max(vi_preds,dim=1)[0].argsort()[:20]
print(torch.max(vi_preds,dim=1)[0][inds])
fig,axs = plt.subplots(figsize=(8,2),nrows=2, ncols=10)
axs = axs.flatten()
for i,im in enumerate(test_loader.dataset.data[inds]):
  ax = axs[i]
  ax.imshow(im,cmap='gray')
  ax.set_axis_off()
  # ax.title.set_text(f'predicted: {vi_preds[i].argmax()}')
# plt.title("Variational Inference in PCA subspace with 10 dimensions")
plt.savefig("vi.pdf")
plt.show()<jupyter_output>tensor([0.2090, 0.2255, 0.2269, 0.2374, 0.2504, 0.2544, 0.2577, 0.2617, 0.2620,
        0.2652, 0.2700, 0.2702, 0.2709, 0.2723, 0.2751, 0.2751, 0.2762, 0.2776,
        0.2778, 0.2793], grad_fn=<IndexBackward>)<jupyter_text>Ensemble<jupyter_code># test_loader = torch.utils.data.DataLoader(test,batch_size=1)
# ensemble_preds = torch.zeros(len(test_loader), 10)

# for i, (x,y) in enumerate(test_loader):
#   x = x.to(device)
#   output = torch.zeros(1,10).to(device)
#   for m in ensemble.ensemble:
#     m.to(device)
#     output += F.softmax(m(x), dim=1)
#   ensemble_preds[i] = output/ensemble.N

inds = torch.max(ensemble_preds,dim=1)[0].argsort()[:20]
print(torch.max(ensemble_preds,dim=1)[0][inds])
fig,axs = plt.subplots(figsize=(5,2),nrows=4, ncols=5)
axs = axs.flatten()
for i,im in enumerate(test_loader.dataset.data[inds]):
  ax = axs[i]
  ax.imshow(im,cmap='gray')
  ax.set_axis_off()
  # ax.title.set_text(f'predicted: {ensemble_preds[i].argmax()}')
plt.savefig("ensemble.png")
plt.show()<jupyter_output><empty_output><jupyter_text>Which gives the 'best' uncertainty?<jupyter_code>def ece(predictions, y, bins=10):
  preds = predictions.argmax(dim=1)
  confidences = predictions.max(dim=1)[0]
  l = np.linspace(0,1,bins+1)
  ece = torch.zeros(bins)
  accs = torch.zeros(bins)
  confs = torch.zeros(bins)
  N = y.size(0)
  for ind, (lower,upper) in enumerate(zip(l[:-1], l[1:])):
    indxs = torch.logical_and((confidences<upper),(confidences>=lower))
    acc = (1.*(preds[indxs] == y[indxs])).mean().item()
    conf = confidences[indxs].mean().item()
    ece[ind] = ((1.*indxs).sum().item()/N) * np.abs(acc-conf)
    accs[ind], confs[ind] = acc, conf
  return ece.nansum(), accs, confs

def plot_ece(confs, label):
  width=0.1
  fig, ax = plt.subplots(figsize=(5,5))
  ax.bar(np.linspace(0,.9,10)+.05, confs, width=width, label=label)
  ax.plot(np.linspace(0,1,11),np.linspace(0,1,11),linestyle='--',label='Perfect uncert.', color='red')
  ax.set_ylabel("Accuracy")
  ax.set_xlabel("Confidence")
  plt.legend()
  plt.savefig(label.replace(" ","-") + ".pdf")
  plt.show()
plt.style.use('classic')
ece_score, accs, conf = ece(vi_preds, test.targets[:,0])
plot_ece(conf, label=f'PCA subspace')
print(ece_score)
print(accs)
print(conf)
'''
ece_score,accs, conf = ece(full_vi_preds, test.targets[:,0])
print(ece_score)
plot_ece(conf, label=f'Full VI subspace')
ece_score, accs, conf = ece(ensemble_preds, test.targets[:,0])
print(ece_score)
plot_ece(conf, label=f'Ensemble')
'''<jupyter_output><empty_output><jupyter_text>_References_ - A very useful repository for a lot of Bayesian NN implementations: https://github.com/JavierAntoran/Bayesian-Neural-Networks- The code for the paper is found at https://github.com/wjmaddox/drbayes and specifically the following notebook was adapted for this demonstration (https://github.com/wjmaddox/drbayes/blob/master/experiments/synthetic_regression/visualizing_uncertainty.ipynb)<jupyter_code>!git config --global user.email "<EMAIL>"
!git config --global user.name "ymohamedahmed"
!git --git-dir=drbayes/.git add -f drbayes/*.pt
!git --git-dir=drbayes/.git commit -m "updated checkpoints"
!git --git-dir=drbayes/.git push
# !cd ../
!git --git-dir=drbayes/.git push
del preds
torch.cuda.empty_cache()
from google.colab import drive
drive.mount('/content/drive')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14536839936826204, 'fraction_numerical': 0.024239511089747992, 'mean_word_length': 5.070862859524802, 'pattern_counts': {'":': 2, '<': 49, '<?xml version=': 0, '>': 52, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '607680', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 5807, 'n_tokens_neox': 5539, 'n_words': 1001, 'path': 'SIBDL_demo_group_13.ipynb', 'repo_name': 'ymohamedahmed/drbayes', 'size': '14609'}