GP_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy import ndimage\n",
    "import os, sys\n",
    "import math\n",
    "import pickle\n",
    "import notebook_utils as nbutils\n",
    "import data_utils as datutil\n",
    "import datetime as dt\n",
    "import hmc\n",
    "from models import *\n",
    "import gpytorch\n",
    "from notebook_utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def learning_rate_mod_factor(epoch, lr_init, lr_end, end_epoch):\n",
    "    lr_ratio = lr_end / lr_init\n",
    "    t = (epoch) / (end_epoch*1.0)\n",
    "    if t < 0.2:\n",
    "        factor = 1.0\n",
    "    elif t <= 0.9:\n",
    "        factor = 1.0 - (1.0 - lr_ratio) * (t - 0.4) / 0.5\n",
    "    else:\n",
    "        factor = lr_ratio\n",
    "    return factor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GP_type1(gpytorch.models.AdditiveGridInducingVariationalGP):\n",
    "    '''\n",
    "    Gaussian Process layer with inducing point approximation of GP\n",
    "\n",
    "    num_dim: incoming feature size\n",
    "    grid_bounds: value bounds for inducing point features\n",
    "    grid_size: number of inducing points to use\n",
    "    '''\n",
    "    def __init__(self, num_dim, grid_bounds=(-10., 10.), grid_size=64):\n",
    "        super(GP_type1, self).__init__(grid_size=grid_size, grid_bounds=[grid_bounds],\n",
    "                                                   num_dim=num_dim, mixing_params=False, sum_output=False)\n",
    "        self.covar_module = gpytorch.kernels.ScaleKernel(\n",
    "            # change the kernel to something else at this point\n",
    "            gpytorch.kernels.RBFKernel(ard_num_dims=num_dim,  # this option sets dimension-wise lengthscale\n",
    "                lengthscale_prior=gpytorch.priors.SmoothedBoxPrior(\n",
    "                    math.exp(-1), math.exp(1), sigma=0.1, transform=torch.exp\n",
    "                )\n",
    "            )\n",
    "        )\n",
    "        self.mean_module = gpytorch.means.ConstantMean()\n",
    "        self.grid_bounds = grid_bounds\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = self.mean_module(x)\n",
    "        covar = self.covar_module(x)\n",
    "        return gpytorch.distributions.MultivariateNormal(mean, covar)\n",
    "\n",
    "\n",
    "class GP_type2(gpytorch.models.AbstractVariationalGP):\n",
    "    '''\n",
    "    Gaussian Process layer with inducing point approximation of GP\n",
    "\n",
    "    num_dim: incoming feature size\n",
    "    grid_bounds: value bounds for inducing point features\n",
    "    grid_size: number of inducing points to use\n",
    "    '''\n",
    "    def __init__(self, num_dim, grid_bounds=(-10., 10.), grid_size=64):\n",
    "        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(\n",
    "            num_inducing_points=grid_size, batch_size=num_dim\n",
    "        )\n",
    "        variational_strategy = gpytorch.variational.AdditiveGridInterpolationVariationalStrategy(\n",
    "            self, grid_size=grid_size, grid_bounds=[grid_bounds], num_dim=num_dim,\n",
    "            variational_distribution=variational_distribution, mixing_params=False, sum_output=False\n",
    "        )\n",
    "        super().__init__(variational_strategy)\n",
    "\n",
    "        self.covar_module = gpytorch.kernels.ScaleKernel(\n",
    "            gpytorch.kernels.RBFKernel(ard_num_dims=num_dim,\n",
    "                lengthscale_prior=gpytorch.priors.SmoothedBoxPrior(\n",
    "                    math.exp(-1), math.exp(1), sigma=0.1, transform=torch.exp\n",
    "                )\n",
    "            )\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPNet(gpytorch.Module):\n",
    "\n",
    "    def __init__(self, device, feature_extractor, num_classes=10, grid_bounds=(-10., 10.), grid_size=64):\n",
    "        '''\n",
    "        Wrapper class for sequentializing a feature extractor and a GP layer\n",
    "        \n",
    "        device: cuda device to place everything into\n",
    "        feature_extractor: base feature extractor, should be set to Identity()\n",
    "                if features are encoded\n",
    "        num_classes: number of classes in classification task\n",
    "        grid_bounds, grid_size: refer to input for GP layer\n",
    "        '''\n",
    "        super(GPNet, self).__init__()\n",
    "        self.feature_extractor = feature_extractor\n",
    "        self.num_dim = feature_extractor.fc_architecture[-1]\n",
    "        self.gp_layer = GP_type1(num_dim=self.num_dim, grid_bounds=grid_bounds, grid_size=grid_size)\n",
    "        self.grid_bounds = grid_bounds\n",
    "        self.device = device\n",
    "        self.feature_extractor.to(device)\n",
    "        self.gp_layer.to(device)\n",
    "\n",
    "    def forward(self, x):\n",
    "        features = self.feature_extractor(x)\n",
    "        features = gpytorch.utils.grid.scale_to_bounds(features, self.grid_bounds[0], self.grid_bounds[1])\n",
    "        res = self.gp_layer(features)\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Identity(nn.Module):\n",
    "    '''\n",
    "    A dummy empty class to place whenever we\n",
    "    do not need any nn block but have to put something\n",
    "    '''\n",
    "    def __init__(self):\n",
    "        super(Identity, self).__init__()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return x\n",
    "\n",
    "class feature(nn.Module):\n",
    "    '''\n",
    "    Wrapper class for feature extractor then a series of dense layers.\n",
    "\n",
    "    base_feature: base feature extractor (maybe resnet till before dense)\n",
    "                set it to Identity class when you have encoded feature.\n",
    "                Otherwise feature extractor parameters will be jointly learned\n",
    "    fc_layers: array containing dense layer lengths starting with base feature dim\n",
    "                e.g. [256, 100, 10] will expect a 256 dimensional input and then\n",
    "                place linear(256, 100) and then linear(100, 10) sequentially\n",
    "    device: cuda device in which (local as well) parameters will be put.\n",
    "    '''\n",
    "    def __init__(self, base_feature, fc_layers=[]):\n",
    "        super(feature, self).__init__()\n",
    "        self.base_layer = base_feature\n",
    "        self.fc_architecture = fc_layers\n",
    "        assert len(fc_layers) > 0, 'FC layer can not be empty, at minimum it needs to be [input_feature_size]'\n",
    "        linear_list = [Identity()]\n",
    "        for comp_idx in range(2*len(fc_layers)-3):\n",
    "            if comp_idx%2==0:\n",
    "                idx = comp_idx // 2\n",
    "                linear_list.append(nn.Linear(fc_layers[idx], fc_layers[idx+1]))\n",
    "            else:\n",
    "                linear_list.append(nn.ReLU())\n",
    "\n",
    "        self.fc_list = nn.Sequential(*linear_list)\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base_layer(x)\n",
    "        if len(self.fc_architecture) > 0:\n",
    "            x = self.fc_list(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data loader initialization\n",
    "trainloader = datutil.generate_dataloaders('ENCODED256_D164_CIFAR10_TRAIN', batch_size=300, shuffle=False, num_workers=2)\n",
    "testloader = datutil.generate_dataloaders('ENCODED256_D164_CIFAR10_TEST', batch_size=200, shuffle=False, num_workers=2)\n",
    "# trainloader = datutil.generate_dataloaders('Fast_resnet30000MixupCutout_train', batch_size=300, shuffle=False, num_workers=2)\n",
    "# validloader = datutil.generate_dataloaders('Fast_resnet30000MixupCutout_train', batch_size=300, \n",
    "#                                            shuffle=False, num_workers=2, start=45000)\n",
    "# testloader = datutil.generate_dataloaders('Fast_resnet30000MixupCutout_test', batch_size=200, shuffle=False, num_workers=2)\n",
    "# trainloader = datutil.generate_dataloaders('Encoded_DR_train', batch_size=300, shuffle=False, num_workers=2)\n",
    "# testloader = datutil.generate_dataloaders('Encoded_DR_test', batch_size=200, shuffle=False, num_workers=2)\n",
    "\n",
    "device = torch.device('cuda:1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPNet(\n",
      "  (feature_extractor): feature(\n",
      "    (base_layer): Identity()\n",
      "    (fc_list): Sequential(\n",
      "      (0): Identity()\n",
      "    )\n",
      "  )\n",
      "  (gp_layer): GP_type1(\n",
      "    (variational_strategy): AdditiveGridInterpolationVariationalStrategy(\n",
      "      (variational_distribution): CholeskyVariationalDistribution()\n",
      "    )\n",
      "    (covar_module): ScaleKernel(\n",
      "      (base_kernel): RBFKernel(\n",
      "        (lengthscale_prior): SmoothedBoxPrior()\n",
      "        (raw_lengthscale_constraint): Positive()\n",
      "      )\n",
      "      (raw_outputscale_constraint): Positive()\n",
      "    )\n",
      "    (mean_module): ConstantMean()\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "num_classes = 10\n",
    "fc_layer_setup = [256] #[256, 160]\n",
    "weight_decay = 0.0003\n",
    "gp_weight_decay = 0.0003\n",
    "lr = 0.1\n",
    "lr_end = 0.0001\n",
    "grid_size = 64\n",
    "\n",
    "# base_model = PreResNet(num_classes=10, depth=164)\n",
    "# base_model.fc = Identity()\n",
    "base_model = Identity()\n",
    "feature_extractor = feature(base_model, fc_layer_setup)\n",
    "final_model = GPNet(device=device, feature_extractor=feature_extractor, num_classes=num_classes, grid_size=grid_size)\n",
    "print(final_model)\n",
    "\n",
    "# likelihood, marginal log-likelihood for GP layers\n",
    "likelihood = gpytorch.likelihoods.SoftmaxLikelihood(fc_layer_setup[-1], num_classes)\n",
    "likelihood.to(device)\n",
    "mll = gpytorch.mlls.VariationalELBO(likelihood, final_model.gp_layer, num_data=len(trainloader.dataset))\n",
    "optimizer = torch.optim.SGD([\n",
    "            {'params': final_model.feature_extractor.parameters(), 'weight_decay': weight_decay},\n",
    "            {'params': final_model.gp_layer.hyperparameters(), 'lr': lr * 0.01, 'weight_decay': gp_weight_decay},\n",
    "            {'params': final_model.gp_layer.variational_parameters(), 'weight_decay': gp_weight_decay},\n",
    "            {'params': likelihood.parameters()}], lr=lr, momentum=0.9)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learning rate for param 0 is currently 0.1000\n",
      "Learning rate for param 1 is currently 0.1000\n",
      "Learning rate for param 2 is currently 0.1000\n",
      "Learning rate for param 3 is currently 0.1000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,     0] loss: 2.308\n",
      "=== Accuracy using SGD params ===\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n",
      "/pytorch/torch/csrc/autograd/python_function.cpp:622: UserWarning: Legacy autograd function with non-static forward method is deprecated and will be removed in 1.3. Please use new-style autograd function with static forward method. (Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy statistics\n",
      "Overall accuracy : 10.3 %\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-feb851288ef6>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"=== Accuracy using SGD params ===\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m     \u001b[0maccuracy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mece\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msce\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnbutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalidate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfinal_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlikelihood\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlikelihood\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdataloader\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtestloader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/lab_work/decoupled-DNN-calibration/notebook_utils.py\u001b[0m in \u001b[0;36mvalidate\u001b[0;34m(model, **kwargs)\u001b[0m\n\u001b[1;32m    109\u001b[0m     \u001b[0mbins\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;36m30.0\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mp\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m30\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    110\u001b[0m     ece_mid, ece_avg, sce_score = calculate_ECE(all_probs=all_prob_list, targets=target_list,\n\u001b[0;32m--> 111\u001b[0;31m                                                 probs=prob_list, preds=pred_list, ECE_bin=bins)\n\u001b[0m\u001b[1;32m    112\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    113\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'ECE values are %.3f, %.3f when mid bin and avg used respectively'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mece_mid\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mece_avg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/lab_work/decoupled-DNN-calibration/notebook_utils.py\u001b[0m in \u001b[0;36mcalculate_ECE\u001b[0;34m(all_probs, targets, probs, preds, ECE_bin)\u001b[0m\n\u001b[1;32m    178\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mclass_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_classes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    179\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mbin_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mECE_bin\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 180\u001b[0;31m                 \u001b[0;32mif\u001b[0m \u001b[0mall_probs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mclass_\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0mECE_bin\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbin_\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    181\u001b[0m                     \u001b[0mSCE_bin_correct\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mclass_\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbin_\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclass_\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mpreds\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    182\u001b[0m                     \u001b[0mSCE_bin_total\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mclass_\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbin_\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "running_loss = 0\n",
    "epoch_length = 50\n",
    "\n",
    "for epoch in range(0, epoch_length):  # loop over the dataset multiple times\n",
    "\n",
    "    factor = learning_rate_mod_factor(epoch, lr, lr_end, epoch_length)\n",
    "    for i, g in enumerate(optimizer.param_groups):\n",
    "        print(\"Learning rate for param %d is currently %.4f\" %(i, lr * factor))\n",
    "        g['lr'] = lr * factor\n",
    "\n",
    "    for i, data in enumerate(trainloader):\n",
    "\n",
    "        inputs, labels = data\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = final_model(inputs)\n",
    "        loss = -mll(output, labels)\n",
    "        loss.sum().backward()\n",
    "        optimizer.step()\n",
    "        running_loss = 0.9*running_loss + 0.1*loss.item() if running_loss != 0 else loss.item()\n",
    "\n",
    "        if i% (len(trainloader) // 1) == 0:\n",
    "            print('[%d, %5d] loss: %.3f' %(epoch + 1, i, running_loss))\n",
    "\n",
    "    print(\"=== Accuracy using SGD params ===\")\n",
    "    accuracy, ece, sce = nbutils.validate(model=final_model, likelihood=likelihood, dataloader=testloader, device=device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "savefile = 'GP_notebook_model_file'\n",
    "savedir = 'saved_models/'\n",
    "checkpoint = {'model_state': final_model.state_dict(),\n",
    "              'likelihood_state' : likelihood.state_dict(),\n",
    "              'optim_state': optimizer.state_dict(),\n",
    "              'acc': 1}\n",
    "curtime = dt.datetime.now()\n",
    "tm = curtime.strftime(\"%Y-%m-%d-%H.%M\")\n",
    "torch.save(checkpoint, savedir + savefile + '-' + tm + '.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(trainloader.dataset))\n",
    "accuracy, ece, sce = nbutils.validate(model=final_model, likelihood=likelihood, dataloader=testloader, device=device)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}