TempScaling.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 *\n",
    "from temp_scaling import _ECELoss\n",
    "import temp_scaling as ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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, device, fc_layers=[]):\n",
    "        super(feature, self).__init__()\n",
    "        self.base_layer = base_feature\n",
    "        self.device = device\n",
    "        self.fc_architecture = fc_layers\n",
    "        if len(fc_layers) > 0:\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], device))\n",
    "                else:\n",
    "                    linear_list.append(nn.ReLU())\n",
    "\n",
    "            self.fc_list = nn.Sequential(*linear_list)\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\n",
    "    \n",
    "    def infer(self, x, num_sample=20):\n",
    "        '''\n",
    "        function to generate class probabilities with\n",
    "        multiple samples from posterior\n",
    "        \n",
    "        x: input (image/encoded features)\n",
    "        num_sample: how many samples to get from posterior\n",
    "        \n",
    "        return: class probabilities of shape (num_sample, x.shape[0], num_classes)\n",
    "        '''\n",
    "        x = self.base_layer(x)\n",
    "        class_prob = torch.zeros((num_sample, x.size()[0], self.fc_architecture[-1]), device=self.device)\n",
    "        for count in range(num_sample):\n",
    "            class_outp = self.fc_list(x)\n",
    "            class_prob[count,:,:] = F.softmax(class_outp, dim=1)\n",
    "\n",
    "        return class_prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data loader initialization\n",
    "trainloader = datutil.generate_dataloaders('Fast_resnet30000MixupCutout_train', batch_size=300, shuffle=False, \n",
    "                                           num_workers=2, end=30000)\n",
    "validloader = datutil.generate_dataloaders('Fast_resnet30000MixupCutout_train', batch_size=500, \n",
    "                                          shuffle=False, num_workers=2, start=30000, end=35000)\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",
    "# 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",
    "device = torch.device('cuda:2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class label_smooth(nn.Module):\n",
    "    def __init__(self, epsilon, num_classes, device):\n",
    "        super(label_smooth, self).__init__()\n",
    "        self.eps = epsilon\n",
    "        self.device = device\n",
    "        self.numc = num_classes\n",
    "\n",
    "    def forward(self, inputs, labels):\n",
    "        probs = F.softmax(inputs, dim=1)\n",
    "        soft_labels = (self.eps / (self.numc - 1)) * torch.ones(probs.shape, device=self.device)\n",
    "        index = labels.unsqueeze(1).type(torch.LongTensor).to(device)\n",
    "        soft_labels.scatter_(dim=1, index=index, value=1-self.eps)\n",
    "        loss = -torch.sum(torch.log(probs) * soft_labels)\n",
    "        return loss / inputs.size()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "feature(\n",
       "  (base_layer): Identity()\n",
       "  (fc_list): Sequential(\n",
       "    (0): Identity()\n",
       "    (1): Linear(in_features=128, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_classes = 10\n",
    "fc_layer_setup = [128, num_classes]\n",
    "weight_decay = 0.0005\n",
    "\n",
    "# base_model = PreResNet(num_classes=10, depth=164)\n",
    "# if customized linear layers are being placed make sure to\n",
    "# remove the already present linear layer from the base feature\n",
    "# base_model.fc = Identity()\n",
    "base_model = Identity()\n",
    "final_model = feature(base_model, device, fc_layer_setup)\n",
    "final_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch_count = 10\n",
    "lr = 0.1\n",
    "end_lr = 0.0001\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "# criterion = limiting_ECE_loss()\n",
    "# criterion = label_smooth(1e-3, 10, device)\n",
    "criterion.to(device)\n",
    "optimizer = torch.optim.SGD(final_model.parameters(), weight_decay=weight_decay, lr=lr, momentum=0.9)#, nesterov=True)\n",
    "# scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, \n",
    "#                                             max_lr = 0.1, \n",
    "#                                             epochs=epoch_count, \n",
    "#                                             steps_per_epoch=len(trainloader), \n",
    "#                                             pct_start=0.25, \n",
    "#                                             anneal_strategy='linear', \n",
    "#                                             cycle_momentum=False, \n",
    "#                                             #cycle_momentum=True, base_momentum=0.9, max_momentum=0.9, \n",
    "#                                             div_factor=25.0, \n",
    "#                                             final_div_factor=10000.0, \n",
    "#                                             last_epoch=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learning rate for param 0 is being set to 0.1000\n",
      "[1,     0] loss: 2.6478\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 70.0 %\n",
      "ECE values are 0.192, 0.198 when mid bin and avg used respectively\n",
      "SCE values are 0.01851\n",
      "Pre-averaging loss: tensor(1.0713, device='cuda:2') Post-averaging loss: tensor(1.0713, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.1000\n",
      "[2,     0] loss: 0.4243\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 91.1 %\n",
      "ECE values are 0.022, 0.020 when mid bin and avg used respectively\n",
      "SCE values are 0.01154\n",
      "Pre-averaging loss: tensor(0.2793, device='cuda:2') Post-averaging loss: tensor(0.2793, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.1400\n",
      "[3,     0] loss: 0.2686\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 90.9 %\n",
      "ECE values are 0.021, 0.017 when mid bin and avg used respectively\n",
      "SCE values are 0.01264\n",
      "Pre-averaging loss: tensor(0.2733, device='cuda:2') Post-averaging loss: tensor(0.2733, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.1200\n",
      "[4,     0] loss: 0.2549\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 90.4 %\n",
      "ECE values are 0.023, 0.020 when mid bin and avg used respectively\n",
      "SCE values are 0.01279\n",
      "Pre-averaging loss: tensor(0.2821, device='cuda:2') Post-averaging loss: tensor(0.2821, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.1000\n",
      "[5,     0] loss: 0.2568\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 90.6 %\n",
      "ECE values are 0.023, 0.020 when mid bin and avg used respectively\n",
      "SCE values are 0.01261\n",
      "Pre-averaging loss: tensor(0.2780, device='cuda:2') Post-averaging loss: tensor(0.2780, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.0800\n",
      "[6,     0] loss: 0.2585\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 90.9 %\n",
      "ECE values are 0.021, 0.017 when mid bin and avg used respectively\n",
      "SCE values are 0.01239\n",
      "Pre-averaging loss: tensor(0.2742, device='cuda:2') Post-averaging loss: tensor(0.2742, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.0600\n",
      "[7,     0] loss: 0.2519\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 91.1 %\n",
      "ECE values are 0.020, 0.016 when mid bin and avg used respectively\n",
      "SCE values are 0.01245\n",
      "Pre-averaging loss: tensor(0.2664, device='cuda:2') Post-averaging loss: tensor(0.2664, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.0401\n",
      "[8,     0] loss: 0.2452\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 91.2 %\n",
      "ECE values are 0.020, 0.016 when mid bin and avg used respectively\n",
      "SCE values are 0.01243\n",
      "Pre-averaging loss: tensor(0.2632, device='cuda:2') Post-averaging loss: tensor(0.2632, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.0201\n",
      "[9,     0] loss: 0.2395\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 91.5 %\n",
      "ECE values are 0.019, 0.014 when mid bin and avg used respectively\n",
      "SCE values are 0.01267\n",
      "Pre-averaging loss: tensor(0.2597, device='cuda:2') Post-averaging loss: tensor(0.2597, device='cuda:2')\n",
      "Learning rate for param 0 is being set to 0.0001\n",
      "[10,     0] loss: 0.2368\n",
      "=== Accuracy using SGD params ===\n",
      "Accuracy statistics\n",
      "Overall accuracy : 91.5 %\n",
      "ECE values are 0.019, 0.015 when mid bin and avg used respectively\n",
      "SCE values are 0.01263\n",
      "Pre-averaging loss: tensor(0.2594, device='cuda:2') Post-averaging loss: tensor(0.2594, device='cuda:2')\n",
      "Time for training 4.3 mins\n"
     ]
    }
   ],
   "source": [
    "running_loss = 0\n",
    "\n",
    "# lr_scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.001, max_lr=0.08, step_size_up=5, step_size_down=10)\n",
    "import time\n",
    "start_t = time.time()\n",
    "for epoch in range(0, epoch_count):  # loop over the dataset multiple times\n",
    "\n",
    "    factor = learning_rate_mod_factor(epoch, lr, end_lr, epoch_count)\n",
    "    for i, g in enumerate(optimizer.param_groups):\n",
    "        print(\"Learning rate for param %d is being set to %.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",
    "        outputs = final_model(inputs)\n",
    "        loss = criterion(outputs, 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: %.4f' %(epoch + 1, i, running_loss))\n",
    "    \n",
    "#     scheduler.step()\n",
    "    print(\"=== Accuracy using SGD params ===\")\n",
    "    accuracy, ece, sce = nbutils.validate(model=final_model, dataloader=testloader, device=device)\n",
    "#     if accuracy >= 93:\n",
    "#         print('93% achieved!')\n",
    "#         break\n",
    "end_t = time.time()\n",
    "print('Time for training %.1f mins' %((end_t - start_t)/60.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "savefile = 'WholeTrainingDataDNN_FullNetworkW2'\n",
    "savedir = 'saved_models/'\n",
    "checkpoint = {'model_state': final_model.state_dict(),\n",
    "              'optim_state': optimizer.state_dict(),\n",
    "              'acc': accuracy}\n",
    "torch.save(checkpoint, savedir + savefile + '.model')\n",
    "\n",
    "# checkpoint = torch.load(savedir + 'interim_coef.model', \n",
    "#                         map_location=device)\n",
    "# final_model.load_state_dict(checkpoint['model_state'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainloader_encoder = datutil.generate_dataloaders('CIFAR10_TRAIN', batch_size=200, shuffle=False, num_workers=2)\n",
    "testloader_encoder = datutil.generate_dataloaders('CIFAR10_TEST', batch_size=200, shuffle=False, num_workers=2)\n",
    "nbutils.encode_dump(final_model.base_layer, '48000CIFAR10ResNet164.pt', trainloader_encoder, device, False)\n",
    "nbutils.encode_dump(final_model.base_layer, '48000CIFAR10ResNet164_test', testloader_encoder, device, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before temperature - NLL: 0.25943, ECE: 0.01211\n",
      "Best loss is: 0.2570318579673767\n",
      "Optimal temperature: 1.13277\n",
      "After temperature - NLL: 0.25703, ECE: 0.00795\n"
     ]
    }
   ],
   "source": [
    "# temperature scale the original model using validation dataset\n",
    "scaled_model = ts.ModelWithTemperature(final_model, device)\n",
    "_ = scaled_model.set_temperature(testloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy statistics\n",
      "Overall accuracy : 91.5 %\n",
      "ECE values are 0.019, 0.015 when mid bin and avg used respectively\n",
      "SCE values are 0.01263\n",
      "Pre-averaging loss: tensor(0.2594, device='cuda:2') Post-averaging loss: tensor(0.2594, device='cuda:2')\n",
      "Accuracy statistics\n",
      "Overall accuracy : 91.5 %\n",
      "ECE values are 0.016, 0.008 when mid bin and avg used respectively\n",
      "SCE values are 0.01456\n",
      "Pre-averaging loss: tensor(0.2570, device='cuda:2') Post-averaging loss: tensor(0.2570, device='cuda:2')\n"
     ]
    }
   ],
   "source": [
    "# calculate accuracy and ECE on the test dataset\n",
    "accuracy, ece, sce = nbutils.validate(model=final_model, dataloader=testloader, device=device)\n",
    "accuracy, ece, sce = nbutils.validate(model=scaled_model, dataloader=testloader, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'temp_scaling' from '/home/rahul/lab_work/decoupled-DNN-calibration/temp_scaling.py'>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import imp\n",
    "imp.reload(ts)"
   ]
  }
 ],
 "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": 1
}