ImageNette-ResNet18.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K9ILC73dI9G9"
   },
   "source": [
    "# Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "is8dFihPJRmW"
   },
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.282051Z",
     "start_time": "2021-12-14T05:11:31.928953Z"
    },
    "executionInfo": {
     "elapsed": 1256,
     "status": "ok",
     "timestamp": 1639339325481,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "OYR89FOfI37J"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.5.3'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# Fast AI (PyTorch wrapper)\n",
    "from fastai import *\n",
    "from fastai.vision.all import *\n",
    "import fastai\n",
    "fastai.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.287324Z",
     "start_time": "2021-12-14T05:11:34.284369Z"
    },
    "executionInfo": {
     "elapsed": 7,
     "status": "ok",
     "timestamp": 1639339325484,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "PW9VU5jIJOeK"
   },
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.294575Z",
     "start_time": "2021-12-14T05:11:34.289698Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "executionInfo": {
     "elapsed": 6,
     "status": "ok",
     "timestamp": 1639339325484,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "EAFb1zouJO0h",
    "outputId": "4224118a-ac2f-4080-a8f1-190b7df4f375"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Tesla V100-SXM2-16GB'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# make sure GPU is being used \n",
    "torch.cuda.current_device() \n",
    "torch.cuda.device(0)\n",
    "torch.cuda.get_device_name(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.323465Z",
     "start_time": "2021-12-14T05:11:34.297560Z"
    }
   },
   "outputs": [],
   "source": [
    "# Notebook auto reloads code. (Ref: http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython)\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Jn0LWLskJToR"
   },
   "source": [
    "## Import Created Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.348584Z",
     "start_time": "2021-12-14T05:11:34.325262Z"
    },
    "id": "A9Df0-4YJV1H"
   },
   "outputs": [],
   "source": [
    "from quantization_functions import quant_aware_resnet_model\n",
    "from quantization_functions import post_training_quant_model\n",
    "from quantization_functions import train_loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AMMfGCwTLWsq"
   },
   "source": [
    "## Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.374926Z",
     "start_time": "2021-12-14T05:11:34.350655Z"
    },
    "id": "b4-R5xIcLmdt"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "TEST_BATCH_SIZE = 32\n",
    "N_CLASS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.399603Z",
     "start_time": "2021-12-14T05:11:34.378017Z"
    }
   },
   "outputs": [],
   "source": [
    "# Download Imagenette 320 pixel\n",
    "\n",
    "path = untar_data(URLs.IMAGENETTE_320)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.423688Z",
     "start_time": "2021-12-14T05:11:34.401428Z"
    }
   },
   "outputs": [],
   "source": [
    "imagenet_stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "\n",
    "train_tfms = transforms.Compose([\n",
    "    transforms.RandomResizedCrop(112),\n",
    "    transforms.RandomHorizontalFlip(), \n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(*imagenet_stats,inplace=True)\n",
    "])\n",
    "\n",
    "test_tfms = transforms.Compose([\n",
    "    transforms.Resize(128),\n",
    "    transforms.CenterCrop(112),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(*imagenet_stats)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.500772Z",
     "start_time": "2021-12-14T05:11:34.425785Z"
    }
   },
   "outputs": [],
   "source": [
    "# PyTorch datasets\n",
    "\n",
    "trainset = datasets.ImageFolder(path/\"train\", train_tfms)\n",
    "testset = datasets.ImageFolder(path/\"val\", test_tfms)\n",
    "\n",
    "# PyTorch data loaders\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n",
    "test_loader = torch.utils.data.DataLoader(testset, batch_size=TEST_BATCH_SIZE, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z378KmkPI8n2"
   },
   "source": [
    "# Resnet 18 Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.524102Z",
     "start_time": "2021-12-14T05:11:34.503441Z"
    },
    "id": "EM6H0tupNLeT"
   },
   "outputs": [],
   "source": [
    "N_EPOCH = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:34.545666Z",
     "start_time": "2021-12-14T05:11:34.526220Z"
    }
   },
   "outputs": [],
   "source": [
    "SAVE_DIR = 'checkpoint/imagenette_resnet18'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XxyL_hxTNL6j"
   },
   "source": [
    "## Base model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:16:15.877208Z",
     "start_time": "2021-12-14T04:16:15.647593Z"
    },
    "id": "Wfw_BQo0fhPF"
   },
   "outputs": [],
   "source": [
    "base_model = torchvision.models.resnet18(pretrained=True)\n",
    "base_model.fc = nn.Linear(base_model.fc.in_features, N_CLASS) # Change top layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:23:57.896737Z",
     "start_time": "2021-12-14T04:16:28.817965Z"
    },
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "id": "_fFaNnSSNXwv",
    "outputId": "afc949b1-8e19-4233-cc01-2074798aa750"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train - epoch:  0: : 74it [01:02,  1.18it/s, loss=0.609]                      \n",
      "val - epoch:  0: : 123it [00:16,  7.63it/s, val_loss=0.244, train_loss=0.609, acc=0.921]                       \n",
      "train - epoch:  1: : 74it [00:27,  2.69it/s, loss=0.291]                      \n",
      "val - epoch:  1: : 123it [00:12,  9.69it/s, val_loss=0.195, train_loss=0.291, acc=0.937]                       \n",
      "train - epoch:  2: : 74it [00:27,  2.65it/s, loss=0.246]                      \n",
      "val - epoch:  2: : 123it [00:13,  9.43it/s, val_loss=0.203, train_loss=0.246, acc=0.942]                       \n",
      "train - epoch:  3: : 74it [00:27,  2.68it/s, loss=0.228]                      \n",
      "val - epoch:  3: : 123it [00:12,  9.52it/s, val_loss=0.202, train_loss=0.228, acc=0.935]                       \n",
      "train - epoch:  4: : 74it [00:28,  2.62it/s, loss=0.209]                      \n",
      "val - epoch:  4: : 123it [00:12,  9.48it/s, val_loss=0.195, train_loss=0.209, acc=0.941]                       \n",
      "train - epoch:  5: : 74it [00:27,  2.72it/s, loss=0.204]                      \n",
      "val - epoch:  5: : 123it [00:12,  9.62it/s, val_loss=0.196, train_loss=0.204, acc=0.938]                       \n",
      "train - epoch:  6: : 74it [00:28,  2.63it/s, loss=0.214]                      \n",
      "val - epoch:  6: : 123it [00:13,  9.25it/s, val_loss=0.167, train_loss=0.214, acc=0.946]                       \n",
      "train - epoch:  7: : 74it [00:27,  2.68it/s, loss=0.193]                      \n",
      "val - epoch:  7: : 123it [00:13,  9.39it/s, val_loss=0.172, train_loss=0.193, acc=0.944]                       \n",
      "train - epoch:  8: : 74it [00:27,  2.72it/s, loss=0.181]                      \n",
      "val - epoch:  8: : 123it [00:13,  9.25it/s, val_loss=0.176, train_loss=0.181, acc=0.945]                       \n",
      "train - epoch:  9: : 74it [00:26,  2.76it/s, loss=0.171]                      \n",
      "val - epoch:  9: : 123it [00:12,  9.51it/s, val_loss=0.204, train_loss=0.171, acc=0.939]                       \n"
     ]
    }
   ],
   "source": [
    "### Train Loop\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(base_model.parameters(), 1e-2, momentum=0.9, weight_decay=1e-5)\n",
    "\n",
    "train_loop.train_model(\n",
    "    train_dl=train_loader, \n",
    "    val_dl=test_loader, \n",
    "    model=base_model, \n",
    "    optimizer=optimizer, \n",
    "    criterion=criterion,\n",
    "    clip_value=1e-2,\n",
    "    epochs=N_EPOCH, save=f\"{SAVE_DIR}/base_model\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:28:55.123535Z",
     "start_time": "2021-12-14T04:28:41.435570Z"
    },
    "colab": {
     "background_save": true
    },
    "id": "aJyLHCM8bBFH",
    "outputId": "a99eb48f-4de7-4dbd-9fda-3a610e181b24"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  9.16it/s, acc=0.946]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9464968152866242\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "base_model = torchvision.models.resnet18(pretrained=False)\n",
    "base_model.fc = nn.Linear(base_model.fc.in_features, N_CLASS) # Change top layer\n",
    "\n",
    "base_model.load_state_dict(torch.load(f'{SAVE_DIR}/base_model/model_weights.pt'))\n",
    "\n",
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, base_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o8y6NZlHQYZH"
   },
   "source": [
    "## Post Training Quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "frTh5sWvmCXA"
   },
   "source": [
    "### 8 bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:26:21.068444Z",
     "start_time": "2021-12-14T04:26:18.452618Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 276,
     "status": "ok",
     "timestamp": 1639008821078,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "GL29AUmrR3Vf",
    "outputId": "edc37568-38c8-45df-94b2-88d16b092703"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys([])\n"
     ]
    }
   ],
   "source": [
    "# Convert base model to a custom quantization layer with the trained weights\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=8, qat=False,\n",
    "                                                  pretrained=f'{SAVE_DIR}/base_model/model_weights.pt')\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:26:49.221023Z",
     "start_time": "2021-12-14T04:26:22.320103Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 109704,
     "status": "ok",
     "timestamp": 1639008945820,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "piIpEik6S1kR",
    "outputId": "a3441a92-e8a4-4dee-f98a-a7362d7495eb"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 74/74 [00:26<00:00,  2.75it/s, acc=0.946]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9455063892702503\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Forward pass to have quantized weights\n",
    "train_loop.test_model(train_loader, c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:26:53.263021Z",
     "start_time": "2021-12-14T04:26:53.213888Z"
    },
    "id": "aoISQcR-PuoY"
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:27:07.763875Z",
     "start_time": "2021-12-14T04:26:54.432639Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 6937,
     "status": "ok",
     "timestamp": 1639008967545,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "K1vVMgjqQbUF",
    "outputId": "0b625bf7-9840-4e65-bef6-315fa33cca8e"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  9.25it/s, acc=0.944]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9439490445859873\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T03:24:47.856767Z",
     "start_time": "2021-12-14T03:24:47.807000Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/ptq8bit_model_weights.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:29:12.889519Z",
     "start_time": "2021-12-14T04:29:10.179969Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys([])\n"
     ]
    }
   ],
   "source": [
    "# Convert base model to a custom quantization layer with the trained weights\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=7, qat=False,\n",
    "                                                  pretrained=f'{SAVE_DIR}/base_model/model_weights.pt')\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:29:41.720337Z",
     "start_time": "2021-12-14T04:29:12.891885Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 74/74 [00:28<00:00,  2.57it/s, acc=0.943]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9432886260428768\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Forward pass to have quantized weights\n",
    "train_loop.test_model(train_loader, c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:29:41.785578Z",
     "start_time": "2021-12-14T04:29:41.725196Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:29:55.752208Z",
     "start_time": "2021-12-14T04:29:41.788220Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  8.83it/s, acc=0.94] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9403821656050956\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:29:55.814558Z",
     "start_time": "2021-12-14T04:29:55.755508Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/ptq7bit_model_weights.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:30:31.247491Z",
     "start_time": "2021-12-14T04:30:28.538656Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys([])\n"
     ]
    }
   ],
   "source": [
    "# Convert base model to a custom quantization layer with the trained weights\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=6, qat=False,\n",
    "                                                  pretrained=f'{SAVE_DIR}/base_model/model_weights.pt')\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:30:59.709373Z",
     "start_time": "2021-12-14T04:30:31.250842Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 74/74 [00:28<00:00,  2.60it/s, acc=0.946]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9455063892702503\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Forward pass to have quantized weights\n",
    "train_loop.test_model(train_loader, c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:30:59.759699Z",
     "start_time": "2021-12-14T04:30:59.711881Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:31:13.355625Z",
     "start_time": "2021-12-14T04:30:59.762274Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  9.07it/s, acc=0.903]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9026751592356688\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:31:13.408929Z",
     "start_time": "2021-12-14T04:31:13.357896Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/ptq6bit_model_weights.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:32:05.140235Z",
     "start_time": "2021-12-14T04:32:02.505991Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys([])\n"
     ]
    }
   ],
   "source": [
    "# Convert base model to a custom quantization layer with the trained weights\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=5, qat=False,\n",
    "                                                  pretrained=f'{SAVE_DIR}/base_model/model_weights.pt')\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:32:34.727075Z",
     "start_time": "2021-12-14T04:32:05.703301Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 74/74 [00:28<00:00,  2.55it/s, acc=0.946]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9459288203611785\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Forward pass to have quantized weights\n",
    "train_loop.test_model(train_loader, c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:32:34.779217Z",
     "start_time": "2021-12-14T04:32:34.730132Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:32:48.598104Z",
     "start_time": "2021-12-14T04:32:34.780919Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  8.92it/s, acc=0.255] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.255031847133758\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:32:48.652636Z",
     "start_time": "2021-12-14T04:32:48.600840Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/ptq5bit_model_weights.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:33:55.983033Z",
     "start_time": "2021-12-14T04:33:53.281090Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys([])\n"
     ]
    }
   ],
   "source": [
    "# Convert base model to a custom quantization layer with the trained weights\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=4, qat=False,\n",
    "                                                  pretrained=f'{SAVE_DIR}/base_model/model_weights.pt')\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:34:23.735834Z",
     "start_time": "2021-12-14T04:33:55.985440Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 74/74 [00:27<00:00,  2.67it/s, acc=0.947]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9465624669975711\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Forward pass to have quantized weights\n",
    "train_loop.test_model(train_loader, c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:34:23.790355Z",
     "start_time": "2021-12-14T04:34:23.739646Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:34:37.484133Z",
     "start_time": "2021-12-14T04:34:23.792931Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  9.00it/s, acc=0.128]  "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.12840764331210192\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:34:37.541890Z",
     "start_time": "2021-12-14T04:34:37.486257Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/ptq4bit_model_weights.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mmpG9_9PO8E5"
   },
   "source": [
    "## Quantization Aware Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4PRbWTsmmRhL"
   },
   "source": [
    "### 8-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:44:51.926963Z",
     "start_time": "2021-12-14T04:44:51.472351Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 872,
     "status": "ok",
     "timestamp": 1639008986864,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "oB6gK6MUO98Z",
    "outputId": "41b3ce18-cba3-4bee-e1f0-8c79f4e48292",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys(['fc.weight', 'fc.bias'])\n"
     ]
    }
   ],
   "source": [
    "# Create model with custom quantization layer from the start\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=8, qat=True, pretrained=True)\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:52:08.939564Z",
     "start_time": "2021-12-14T04:44:51.929441Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 331716,
     "status": "ok",
     "timestamp": 1639009319891,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "QJFoGEjZNTkA",
    "outputId": "9973e40d-0a42-499a-d085-8f2746212eef"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train - epoch:  0: : 74it [00:29,  2.52it/s, loss=0.886]                      \n",
      "val - epoch:  0: : 123it [00:14,  8.22it/s, val_loss=0.282, train_loss=0.886, acc=0.907]                       \n",
      "train - epoch:  1: : 74it [00:28,  2.57it/s, loss=0.393]                      \n",
      "val - epoch:  1: : 123it [00:14,  8.60it/s, val_loss=0.246, train_loss=0.393, acc=0.913]                       \n",
      "train - epoch:  2: : 74it [00:27,  2.69it/s, loss=0.328]                      \n",
      "val - epoch:  2: : 123it [00:14,  8.34it/s, val_loss=0.221, train_loss=0.328, acc=0.929]                       \n",
      "train - epoch:  3: : 74it [00:29,  2.53it/s, loss=0.293]                      \n",
      "val - epoch:  3: : 123it [00:14,  8.62it/s, val_loss=0.21, train_loss=0.293, acc=0.934]                        \n",
      "train - epoch:  4: : 74it [00:29,  2.55it/s, loss=0.283]                      \n",
      "val - epoch:  4: : 123it [00:14,  8.38it/s, val_loss=0.203, train_loss=0.283, acc=0.935]                       \n",
      "train - epoch:  5: : 74it [00:28,  2.60it/s, loss=0.261]                      \n",
      "val - epoch:  5: : 123it [00:14,  8.72it/s, val_loss=0.192, train_loss=0.261, acc=0.937]                       \n",
      "train - epoch:  6: : 74it [00:28,  2.63it/s, loss=0.261]                      \n",
      "val - epoch:  6: : 123it [00:14,  8.68it/s, val_loss=0.178, train_loss=0.261, acc=0.942]                       \n",
      "train - epoch:  7: : 74it [00:28,  2.64it/s, loss=0.244]                      \n",
      "val - epoch:  7: : 123it [00:14,  8.65it/s, val_loss=0.168, train_loss=0.244, acc=0.945]                       \n",
      "train - epoch:  8: : 74it [00:28,  2.61it/s, loss=0.238]                      \n",
      "val - epoch:  8: : 123it [00:14,  8.47it/s, val_loss=0.183, train_loss=0.238, acc=0.941]                       \n",
      "train - epoch:  9: : 74it [00:29,  2.55it/s, loss=0.222]                      \n",
      "val - epoch:  9: : 123it [00:14,  8.37it/s, val_loss=0.18, train_loss=0.222, acc=0.943]                        \n"
     ]
    }
   ],
   "source": [
    "# Training Loop\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(c_base_model.parameters(), 1e-3, momentum=0.9, weight_decay=1e-5)\n",
    "\n",
    "train_loop.train_model(\n",
    "    train_dl=train_loader, \n",
    "    val_dl=test_loader, \n",
    "    model=c_base_model, \n",
    "    optimizer=optimizer, \n",
    "    criterion=criterion,\n",
    "    clip_value=1e-2,\n",
    "    epochs=N_EPOCH, save=f\"{SAVE_DIR}/qat8bit\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:52:08.990592Z",
     "start_time": "2021-12-14T04:52:08.942395Z"
    },
    "id": "muNQF34RWwZ6"
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:52:22.883123Z",
     "start_time": "2021-12-14T04:52:08.992623Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 6681,
     "status": "ok",
     "timestamp": 1639009326731,
     "user": {
      "displayName": "Rifqi Luthfan",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi6REDrzyVNDZl2VT9M3BzwlBYF_1Y_KHDREx_t=s64",
      "userId": "07442021322718158104"
     },
     "user_tz": 300
    },
    "id": "UPl1MlvHJHOY",
    "outputId": "8fd185d3-fcd2-4124-e81e-ec5f23efad93"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:13<00:00,  8.87it/s, acc=0.944]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9439490445859873\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:54:24.960037Z",
     "start_time": "2021-12-14T04:54:24.851917Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/qat8bit/model_weights_quantized.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T04:54:45.682446Z",
     "start_time": "2021-12-14T04:54:45.235263Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys(['fc.weight', 'fc.bias'])\n"
     ]
    }
   ],
   "source": [
    "# Create model with custom quantization layer from the start\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=7, qat=True, pretrained=True)\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:02:05.519634Z",
     "start_time": "2021-12-14T04:54:46.838764Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train - epoch:  0: : 74it [00:29,  2.50it/s, loss=0.982]                      \n",
      "val - epoch:  0: : 123it [00:14,  8.32it/s, val_loss=0.342, train_loss=0.982, acc=0.887]                       \n",
      "train - epoch:  1: : 74it [00:28,  2.58it/s, loss=0.395]                      \n",
      "val - epoch:  1: : 123it [00:14,  8.43it/s, val_loss=0.272, train_loss=0.395, acc=0.911]                       \n",
      "train - epoch:  2: : 74it [00:29,  2.54it/s, loss=0.355]                      \n",
      "val - epoch:  2: : 123it [00:14,  8.42it/s, val_loss=0.23, train_loss=0.355, acc=0.928]                        \n",
      "train - epoch:  3: : 74it [00:28,  2.60it/s, loss=0.331]                      \n",
      "val - epoch:  3: : 123it [00:14,  8.24it/s, val_loss=0.216, train_loss=0.331, acc=0.93]                        \n",
      "train - epoch:  4: : 74it [00:28,  2.61it/s, loss=0.297]                      \n",
      "val - epoch:  4: : 123it [00:14,  8.46it/s, val_loss=0.21, train_loss=0.297, acc=0.933]                        \n",
      "train - epoch:  5: : 74it [00:28,  2.63it/s, loss=0.281]                      \n",
      "val - epoch:  5: : 123it [00:14,  8.39it/s, val_loss=0.202, train_loss=0.281, acc=0.933]                       \n",
      "train - epoch:  6: : 74it [00:28,  2.60it/s, loss=0.276]                      \n",
      "val - epoch:  6: : 123it [00:14,  8.70it/s, val_loss=0.202, train_loss=0.276, acc=0.933]                       \n",
      "train - epoch:  7: : 74it [00:28,  2.59it/s, loss=0.263]                      \n",
      "val - epoch:  7: : 123it [00:14,  8.36it/s, val_loss=0.181, train_loss=0.263, acc=0.939]                       \n",
      "train - epoch:  8: : 74it [00:28,  2.61it/s, loss=0.249]                      \n",
      "val - epoch:  8: : 123it [00:14,  8.56it/s, val_loss=0.187, train_loss=0.249, acc=0.943]                       \n",
      "train - epoch:  9: : 74it [00:28,  2.59it/s, loss=0.251]                      \n",
      "val - epoch:  9: : 123it [00:14,  8.39it/s, val_loss=0.183, train_loss=0.251, acc=0.94]                        \n"
     ]
    }
   ],
   "source": [
    "# Training Loop\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(c_base_model.parameters(), 1e-3, momentum=0.9, weight_decay=1e-5)\n",
    "\n",
    "train_loop.train_model(\n",
    "    train_dl=train_loader, \n",
    "    val_dl=test_loader, \n",
    "    model=c_base_model, \n",
    "    optimizer=optimizer, \n",
    "    criterion=criterion,\n",
    "    clip_value=1e-2,\n",
    "    epochs=N_EPOCH, save=f\"{SAVE_DIR}/qat7bit\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:02:05.563587Z",
     "start_time": "2021-12-14T05:02:05.522110Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:02:19.619758Z",
     "start_time": "2021-12-14T05:02:05.564974Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:14<00:00,  8.77it/s, acc=0.94] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.9401273885350319\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:02:20.828011Z",
     "start_time": "2021-12-14T05:02:20.773944Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/qat7bit/model_weights_quantized.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T03:49:16.247583Z",
     "start_time": "2021-12-14T03:49:15.822492Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys(['fc.weight', 'fc.bias'])\n"
     ]
    }
   ],
   "source": [
    "# Create model with custom quantization layer from the start\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=6, qat=True, pretrained=True)\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T03:54:07.264084Z",
     "start_time": "2021-12-14T03:49:16.250195Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train - epoch:  0: : 196it [00:25,  7.60it/s, loss=1.06]                       \n",
      "val - epoch:  0: : 40it [00:03, 12.22it/s, val_loss=0.847, train_loss=1.06, acc=0.701]                      \n",
      "train - epoch:  1: : 196it [00:26,  7.48it/s, loss=0.728]                       \n",
      "val - epoch:  1: : 40it [00:03, 13.22it/s, val_loss=0.805, train_loss=0.728, acc=0.733]                      \n",
      "train - epoch:  2: : 196it [00:25,  7.59it/s, loss=0.626]                       \n",
      "val - epoch:  2: : 40it [00:02, 13.47it/s, val_loss=0.698, train_loss=0.626, acc=0.77]                       \n",
      "train - epoch:  3: : 196it [00:25,  7.63it/s, loss=0.563]                       \n",
      "val - epoch:  3: : 40it [00:02, 13.60it/s, val_loss=0.707, train_loss=0.563, acc=0.763]                      \n",
      "train - epoch:  4: : 196it [00:25,  7.67it/s, loss=0.515]                       \n",
      "val - epoch:  4: : 40it [00:02, 13.64it/s, val_loss=0.677, train_loss=0.515, acc=0.781]                      \n",
      "train - epoch:  5: : 196it [00:25,  7.70it/s, loss=0.493]                       \n",
      "val - epoch:  5: : 40it [00:03, 13.01it/s, val_loss=0.716, train_loss=0.493, acc=0.78]                       \n",
      "train - epoch:  6: : 196it [00:25,  7.68it/s, loss=0.453]                       \n",
      "val - epoch:  6: : 40it [00:02, 14.36it/s, val_loss=0.782, train_loss=0.453, acc=0.749]                      \n",
      "train - epoch:  7: : 196it [00:25,  7.72it/s, loss=0.455]                       \n",
      "val - epoch:  7: : 40it [00:02, 13.64it/s, val_loss=0.731, train_loss=0.455, acc=0.781]                      \n",
      "train - epoch:  8: : 196it [00:25,  7.67it/s, loss=0.452]                       \n",
      "val - epoch:  8: : 40it [00:02, 13.81it/s, val_loss=0.734, train_loss=0.452, acc=0.776]                      \n",
      "train - epoch:  9: : 196it [00:25,  7.79it/s, loss=0.452]                       \n",
      "val - epoch:  9: : 40it [00:02, 13.39it/s, val_loss=0.755, train_loss=0.452, acc=0.77]                       \n"
     ]
    }
   ],
   "source": [
    "# Training Loop\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(c_base_model.parameters(), 1e-2, momentum=0.9, weight_decay=1e-5)\n",
    "\n",
    "train_loop.train_model(\n",
    "    train_dl=train_loader, \n",
    "    val_dl=test_loader, \n",
    "    model=c_base_model, \n",
    "    optimizer=optimizer, \n",
    "    criterion=criterion,\n",
    "    clip_value=1e-2,\n",
    "    epochs=N_EPOCH, save=f\"{SAVE_DIR}/qat6bit\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T03:54:07.302282Z",
     "start_time": "2021-12-14T03:54:07.265989Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T03:54:11.033678Z",
     "start_time": "2021-12-14T03:54:07.304337Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 40/40 [00:03<00:00, 10.79it/s, acc=0.769]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.7689\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T03:54:11.089964Z",
     "start_time": "2021-12-14T03:54:11.036938Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/qat6bit/model_weights_quantized.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:02:39.195468Z",
     "start_time": "2021-12-14T05:02:38.756042Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys(['fc.weight', 'fc.bias'])\n"
     ]
    }
   ],
   "source": [
    "# Create model with custom quantization layer from the start\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=5, qat=True, pretrained=True)\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:09:56.556218Z",
     "start_time": "2021-12-14T05:02:39.224204Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train - epoch:  0: : 74it [00:28,  2.61it/s, loss=1.9]                       \n",
      "val - epoch:  0: : 123it [00:14,  8.59it/s, val_loss=1.06, train_loss=1.9, acc=0.639]                       \n",
      "train - epoch:  1: : 74it [00:28,  2.62it/s, loss=1.1]                       \n",
      "val - epoch:  1: : 123it [00:14,  8.45it/s, val_loss=0.794, train_loss=1.1, acc=0.735]                       \n",
      "train - epoch:  2: : 74it [00:28,  2.64it/s, loss=0.954]                      \n",
      "val - epoch:  2: : 123it [00:15,  8.15it/s, val_loss=0.692, train_loss=0.954, acc=0.766]                       \n",
      "train - epoch:  3: : 74it [00:29,  2.52it/s, loss=0.885]                      \n",
      "val - epoch:  3: : 123it [00:14,  8.43it/s, val_loss=0.631, train_loss=0.885, acc=0.786]                       \n",
      "train - epoch:  4: : 74it [00:28,  2.63it/s, loss=0.922]                      \n",
      "val - epoch:  4: : 123it [00:14,  8.46it/s, val_loss=0.694, train_loss=0.922, acc=0.779]                       \n",
      "train - epoch:  5: : 74it [00:28,  2.56it/s, loss=0.873]                      \n",
      "val - epoch:  5: : 123it [00:14,  8.51it/s, val_loss=0.625, train_loss=0.873, acc=0.79]                        \n",
      "train - epoch:  6: : 74it [00:28,  2.62it/s, loss=0.87]                       \n",
      "val - epoch:  6: : 123it [00:14,  8.54it/s, val_loss=0.614, train_loss=0.87, acc=0.797]                       \n",
      "train - epoch:  7: : 74it [00:28,  2.58it/s, loss=0.834]                      \n",
      "val - epoch:  7: : 123it [00:14,  8.46it/s, val_loss=0.612, train_loss=0.834, acc=0.798]                       \n",
      "train - epoch:  8: : 74it [00:28,  2.61it/s, loss=0.809]                      \n",
      "val - epoch:  8: : 123it [00:14,  8.67it/s, val_loss=0.555, train_loss=0.809, acc=0.817]                       \n",
      "train - epoch:  9: : 74it [00:28,  2.55it/s, loss=0.758]                      \n",
      "val - epoch:  9: : 123it [00:14,  8.25it/s, val_loss=0.522, train_loss=0.758, acc=0.829]                       \n"
     ]
    }
   ],
   "source": [
    "# Training Loop\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(c_base_model.parameters(), 1e-3, momentum=0.9, weight_decay=1e-5)\n",
    "\n",
    "train_loop.train_model(\n",
    "    train_dl=train_loader, \n",
    "    val_dl=test_loader, \n",
    "    model=c_base_model, \n",
    "    optimizer=optimizer, \n",
    "    criterion=criterion,\n",
    "    clip_value=1e-2,\n",
    "    epochs=N_EPOCH, save=f\"{SAVE_DIR}/qat5bit\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:09:56.600091Z",
     "start_time": "2021-12-14T05:09:56.558869Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:10:12.034146Z",
     "start_time": "2021-12-14T05:09:56.601713Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:15<00:00,  7.99it/s, acc=0.806]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.8061146496815287\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:10:12.128542Z",
     "start_time": "2021-12-14T05:10:12.041279Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/qat5bit/model_weights_quantized.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4-bit quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:11:38.958479Z",
     "start_time": "2021-12-14T05:11:38.486016Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "remained state dict odict_keys(['fc.weight', 'fc.bias'])\n"
     ]
    }
   ],
   "source": [
    "# Create model with custom quantization layer from the start\n",
    "c_base_model = quant_aware_resnet_model.CResnet18(num_class=10, q_num_bit=4, qat=True, pretrained=True)\n",
    "c_base_model.quantize(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:19:07.406202Z",
     "start_time": "2021-12-14T05:11:38.961498Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train - epoch:  0: : 74it [00:30,  2.45it/s, loss=2.54]                      \n",
      "val - epoch:  0: : 123it [00:14,  8.48it/s, val_loss=2.34, train_loss=2.54, acc=0.121]                       \n",
      "train - epoch:  1: : 74it [00:29,  2.51it/s, loss=2.31]                      \n",
      "val - epoch:  1: : 123it [00:14,  8.37it/s, val_loss=2.28, train_loss=2.31, acc=0.145]                       \n",
      "train - epoch:  2: : 74it [00:29,  2.54it/s, loss=2.53]                      \n",
      "val - epoch:  2: : 123it [00:14,  8.40it/s, val_loss=2.37, train_loss=2.53, acc=0.141]                       \n",
      "train - epoch:  3: : 74it [00:28,  2.56it/s, loss=2.36]                      \n",
      "val - epoch:  3: : 123it [00:14,  8.47it/s, val_loss=2.25, train_loss=2.36, acc=0.164]                       \n",
      "train - epoch:  4: : 74it [00:27,  2.65it/s, loss=2.26]                      \n",
      "val - epoch:  4: : 123it [00:14,  8.52it/s, val_loss=2.24, train_loss=2.26, acc=0.15]                        \n",
      "train - epoch:  5: : 74it [00:30,  2.44it/s, loss=2.24]                      \n",
      "val - epoch:  5: : 123it [00:14,  8.65it/s, val_loss=2.23, train_loss=2.24, acc=0.148]                       \n",
      "train - epoch:  6: : 74it [00:28,  2.58it/s, loss=2.24]                      \n",
      "val - epoch:  6: : 123it [00:14,  8.57it/s, val_loss=2.23, train_loss=2.24, acc=0.146]                       \n",
      "train - epoch:  7: : 74it [00:30,  2.42it/s, loss=2.23]                      \n",
      "val - epoch:  7: : 123it [00:15,  8.00it/s, val_loss=2.24, train_loss=2.23, acc=0.139]                       \n",
      "train - epoch:  8: : 74it [00:30,  2.46it/s, loss=2.23]                      \n",
      "val - epoch:  8: : 123it [00:16,  7.64it/s, val_loss=2.24, train_loss=2.23, acc=0.137]                       \n",
      "train - epoch:  9: : 74it [00:30,  2.45it/s, loss=2.23]                      \n",
      "val - epoch:  9: : 123it [00:14,  8.29it/s, val_loss=2.24, train_loss=2.23, acc=0.135]                       \n"
     ]
    }
   ],
   "source": [
    "# Training Loop\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(c_base_model.parameters(), 1e-3, momentum=0.9, weight_decay=1e-5)\n",
    "\n",
    "train_loop.train_model(\n",
    "    train_dl=train_loader, \n",
    "    val_dl=test_loader, \n",
    "    model=c_base_model, \n",
    "    optimizer=optimizer, \n",
    "    criterion=criterion,\n",
    "    clip_value=1e-2,\n",
    "    epochs=N_EPOCH, save=f\"{SAVE_DIR}/qat4bit\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:19:07.452494Z",
     "start_time": "2021-12-14T05:19:07.408916Z"
    }
   },
   "outputs": [],
   "source": [
    "# Convert to quantized model\n",
    "q_base_model = post_training_quant_model.QResnet18(num_class=10)\n",
    "q_base_model.convert_from(c_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:19:21.939897Z",
     "start_time": "2021-12-14T05:19:07.454058Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:14<00:00,  8.51it/s, acc=0.133]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc: 0.1332484076433121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Validation accuracy\n",
    "train_loop.test_model(test_loader, q_base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-14T05:19:22.004891Z",
     "start_time": "2021-12-14T05:19:21.942851Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(f'{SAVE_DIR}/qat4bit/model_weights_quantized.pt', 'wb') as f:\n",
    "    torch.save(q_base_model.state_dict(), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyOh2T6Tj/G/RPl19+iwVusG",
   "collapsed_sections": [],
   "name": "CIFAR10.ipynb",
   "provenance": []
  },
  "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.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "302.4px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}