notebooks/BayesianModels_And_Custom_Pyro_Modules.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0,'..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "import priors\n",
    "from train import train, get_weighted_single_eval_pos_sampler\n",
    "import encoders\n",
    "import positional_encodings\n",
    "import utils\n",
    "import bar_distribution\n",
    "import decoders\n",
    "from datasets import *\n",
    "import os\n",
    "\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pyro.distributions as dist\n",
    "\n",
    "import torch.nn as nn\n",
    "import os.path\n",
    "import glob\n",
    "\n",
    "from mcmc_svi_transformer_on_bayesian import get_model, get_default_model_spec, generate_toy_data, load_results, plot_with_confidence_intervals, training_steps, training_samples, get_default_evaluation_points, compute_mean_and_conf_interval, eval_transformer\n",
    "from pyro.nn import PyroModule, PyroParam, PyroSample\n",
    "import pyro\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## DEFINE A PRIOR MODEL ##\n",
    "# We define a Bayesian Model as a prior for all methods\n",
    "# This can be replaced by other models that inherit from PyroModule.\n",
    "class BayesianModel(PyroModule):\n",
    "    def __init__(self, model_spec, device='cuda'):\n",
    "        super().__init__()\n",
    "\n",
    "        self.device = device\n",
    "        self.num_features = model_spec['num_features']\n",
    "\n",
    "        mu, sigma = torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)\n",
    "\n",
    "        self.fc1 = PyroModule[nn.Linear](self.num_features, model_spec['embed'])\n",
    "        self.fc1.weight = PyroSample(\n",
    "            dist.Normal(mu, sigma).expand([model_spec['embed'], self.num_features]).to_event(2))\n",
    "        self.fc1.bias = PyroSample(dist.Normal(mu, sigma).expand([model_spec['embed']]).to_event(1))\n",
    "\n",
    "        self.fc2 = PyroModule[nn.Linear](model_spec['embed'], 2)\n",
    "        self.fc2.weight = PyroSample(dist.Normal(mu, sigma).expand([2, model_spec['embed']]).to_event(2))\n",
    "        self.fc2.bias = PyroSample(dist.Normal(mu, sigma).expand([2]).to_event(1))\n",
    "\n",
    "        self.model = torch.nn.Sequential(self.fc1, self.fc2)\n",
    "\n",
    "        self.to(self.device)\n",
    "\n",
    "    def forward(self, x=None, y=None, seq_len=1):\n",
    "        if x is None:\n",
    "            with pyro.plate(\"x_plate\", seq_len):\n",
    "                d_ = dist.Normal(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)).expand(\n",
    "                    [self.num_features]).to_event(1)\n",
    "                x = pyro.sample(\"x\", d_)\n",
    "\n",
    "        out = self.model(x)\n",
    "        mu = out.squeeze()\n",
    "        softmax = torch.nn.Softmax(dim=1)\n",
    "        with pyro.plate(\"data\", out.shape[0]):\n",
    "            s = softmax(mu)\n",
    "            obs = pyro.sample('obs', dist.Categorical(probs=s), obs=y).float()\n",
    "\n",
    "        return x, obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_directory = '../results' # Where to save results\n",
    "model_spec_size = 'small' # Size of the BNN model to evaluate, also try big\n",
    "bptt = 100 # Number of samples in each dataset\n",
    "\n",
    "# Training samples seen after which to evaluate the methods\n",
    "evaluation_points = [2, 7, 12, 17, 22, 27, 32, 37, 42, 47, 52, 57, 62, 67, 72, 77, 82, 87, 92]\n",
    "model_spec = get_default_model_spec(model_spec_size)\n",
    "\n",
    "# Function which generates a model from the prior\n",
    "model_sampler = lambda : BayesianModel(model_spec, device = device)\n",
    "\n",
    "global_results = {} # Dict in which to save results\n",
    "task = 'samples' # Task to evaluate, only option is samples, keep fixed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir {results_directory}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Evaluate SVI and MCMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "method = 'svi'\n",
    "steps = 1\n",
    "device = 'cpu'\n",
    "path_interfix = f'{results_directory}/timing_{model_spec_size}_model_test'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_model = model_sampler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "!mkdir {path_interfix}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# This evaluates baseline method (SVI or MCMC) with varying number of trainings samples and steps number of inference steps\n",
    "\n",
    "X, y = generate_toy_data(test_model, bptt, device)\n",
    "\n",
    "training_samples(method\n",
    "                 , X\n",
    "                 , y\n",
    "                 , model_sampler\n",
    "                 , evaluation_points\n",
    "                 , steps=steps\n",
    "                 , device=device\n",
    "                 , path_interfix=path_interfix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Training Transformer on Prior (Skip this step to reuse results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "device = 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "config = {'lr': 2.006434218345026e-05\n",
    " , 'epochs': 160\n",
    " , 'dropout': 0.0\n",
    " , 'emsize': 256\n",
    " , 'batch_size': 256\n",
    " , 'nlayers': 5\n",
    " , 'num_outputs': 1\n",
    " , 'num_features': model_spec['num_features']\n",
    " , 'steps_per_epoch': 100\n",
    " , 'nhead': 4\n",
    " , 'seq_len': model_spec['seq_len']\n",
    " , 'nhid_factor': 2}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "transformer_model = get_model(model_sampler, config, should_train = True)\n",
    "model_path = os.path.join(results_directory, f'bayesian_models_transformer_checkpoint_{model_spec_size}_epochs_'+config['epochs']+'.cpkt')\n",
    "torch.save((transformer_model[2].state_dict(), None), model_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluating Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_epoch = config['epochs']\n",
    "transformer_model = get_model(model_sampler, config, should_train = False)\n",
    "path = os.path.join(results_directory, f'bayesian_models_transformer_checkpoint_{model_spec_size}_epochs_{loaded_epoch}.cpkt')\n",
    "model_state, optimizer_state = torch.load(path)\n",
    "transformer_model[2].load_state_dict(model_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y = generate_toy_data(test_model, bptt, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_acc = []\n",
    "results_nll = []\n",
    "transformer_model[2].eval()\n",
    "for training_samples_n in evaluation_points:\n",
    "    print(training_samples_n)\n",
    "    acc, nll, elapsed = eval_transformer(X, y, model=transformer_model[2], training_samples_n=training_samples_n, device=device)\n",
    "    results_acc.append(acc)\n",
    "    results_nll.append(nll)\n",
    "mean = np.array([compute_mean_and_conf_interval(nll)[0] for nll in results_nll])\n",
    "conf = np.array([compute_mean_and_conf_interval(nll)[1] for nll in results_nll])\n",
    "\n",
    "global_results['transformer'] = (None, np.array(evaluation_points), mean, conf)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files, times, samples, mean, conf = load_results(f'{results_directory}/timing_{model_spec_size}_model/results_svi_training_{task}', task=task)\n",
    "global_results['svi'] = (times/100, samples, mean, conf)\n",
    "files, times, samples, mean, conf = load_results(f'{results_directory}/timing_{model_spec_size}_model/results_mcmc_training_{task}', task=task)\n",
    "global_results['mcmc'] = (times/100, samples,mean, conf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_min = min([global_results[k][2].min() for k in global_results])\n",
    "y_max = max([global_results[k][2].max() for k in global_results])\n",
    "\n",
    "fig2 = plt.figure(constrained_layout=True, figsize=(7, 4))\n",
    "axes = plt.axes()\n",
    "axes.set_xlim(2, 100)\n",
    "#axes.set_ylim(y_min, y_max)\n",
    "for k in global_results:\n",
    "    plot_with_confidence_intervals(plt, global_results[k][1], global_results[k][2], global_results[k][3], label=k)\n",
    "    #plt.plot(global_results_train_steps[k][1], global_results_train_steps[k][0], label=k)\n",
    "plt.legend(loc=\"upper right\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}