{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import collections\n",
    "\n",
    "import torch\n",
    "import torch.utils.data\n",
    "from torch import nn\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "\n",
    "from pytorch_quantization import nn as quant_nn\n",
    "from pytorch_quantization import calib\n",
    "from pytorch_quantization.tensor_quant import QuantDescriptor\n",
    "\n",
    "from absl import logging\n",
    "logging.set_verbosity(logging.FATAL)  # Disable logging as they are too noisy in notebook\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For simplicity, import train and eval functions from the train script from torchvision instead of copything them here\n",
    "# Download torchvision from https://github.com/pytorch/vision\n",
    "sys.path.append(\"/raid/skyw/models/torchvision/references/classification/\")\n",
    "from train import evaluate, train_one_epoch, load_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set default QuantDescriptor to use histogram based calibration for activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "quant_desc_input = QuantDescriptor(calib_method='histogram')\n",
    "quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)\n",
    "quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize quantized modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_quantization import quant_modules\n",
    "quant_modules.initialize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create model with pretrained weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): QuantConv2d(\n",
       "    3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False\n",
       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "  )\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): QuantLinear(\n",
       "    in_features=2048, out_features=1000, bias=True\n",
       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = torchvision.models.resnet50(pretrained=True, progress=False)\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data\n",
      "Loading training data\n",
      "Took 3.580507755279541\n",
      "Loading validation data\n",
      "Creating data loaders\n"
     ]
    }
   ],
   "source": [
    "data_path = \"/raid/data/imagenet/imagenet_pytorch\"\n",
    "batch_size = 512\n",
    "\n",
    "traindir = os.path.join(data_path, 'train')\n",
    "valdir = os.path.join(data_path, 'val')\n",
    "_args = collections.namedtuple('mock_args', ['model', 'distributed', 'cache_dataset'])\n",
    "dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, _args(model='resnet50', distributed=False, cache_dataset=False))\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=batch_size,\n",
    "    sampler=train_sampler, num_workers=4, pin_memory=True)\n",
    "\n",
    "data_loader_test = torch.utils.data.DataLoader(\n",
    "    dataset_test, batch_size=batch_size,\n",
    "    sampler=test_sampler, num_workers=4, pin_memory=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calibrate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_stats(model, data_loader, num_batches):\n",
    "    \"\"\"Feed data to the network and collect statistic\"\"\"\n",
    "\n",
    "    # Enable calibrators\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
    "            if module._calibrator is not None:\n",
    "                module.disable_quant()\n",
    "                module.enable_calib()\n",
    "            else:\n",
    "                module.disable()\n",
    "\n",
    "    for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):\n",
    "        model(image.cuda())\n",
    "        if i >= num_batches:\n",
    "            break\n",
    "\n",
    "    # Disable calibrators\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
    "            if module._calibrator is not None:\n",
    "                module.enable_quant()\n",
    "                module.disable_calib()\n",
    "            else:\n",
    "                module.enable()\n",
    "            \n",
    "def compute_amax(model, **kwargs):\n",
    "    # Load calib result\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
    "            if module._calibrator is not None:\n",
    "                if isinstance(module._calibrator, calib.MaxCalibrator):\n",
    "                    module.load_calib_amax()\n",
    "                else:\n",
    "                    module.load_calib_amax(**kwargs)\n",
    "#             print(F\"{name:40}: {module}\")\n",
    "    model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [04:50<00:00, 111.13s/it]"
     ]
    }
   ],
   "source": [
    "# It is a bit slow since we collect histograms on CPU\n",
    "with torch.no_grad():\n",
    "    collect_stats(model, data_loader, num_batches=2)\n",
    "    compute_amax(model, method=\"percentile\", percentile=99.99)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now evaluate the calibrated model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test:  [ 0/98]  eta: 0:05:53  loss: 0.5656 (0.5656)  acc1: 85.7422 (85.7422)  acc5: 96.0938 (96.0938)  time: 3.6079  data: 2.8152  max mem: 5880\n",
      "Test:  [20/98]  eta: 0:01:07  loss: 0.6741 (0.6825)  acc1: 82.8125 (82.4219)  acc5: 95.8984 (95.7682)  time: 0.7343  data: 0.0002  max mem: 5882\n",
      "Test:  [40/98]  eta: 0:00:46  loss: 0.6995 (0.7157)  acc1: 80.0781 (81.4024)  acc5: 96.0938 (95.7412)  time: 0.7226  data: 0.0002  max mem: 5882\n",
      "Test:  [60/98]  eta: 0:00:29  loss: 1.1064 (0.8590)  acc1: 71.4844 (78.2627)  acc5: 91.0156 (94.1150)  time: 0.7259  data: 0.0002  max mem: 5882\n",
      "Test:  [80/98]  eta: 0:00:13  loss: 1.1220 (0.9372)  acc1: 72.4609 (76.7072)  acc5: 89.6484 (93.1375)  time: 0.7220  data: 0.0002  max mem: 5882\n",
      "Test: Total time: 0:01:13\n",
      " * Acc@1 76.138 Acc@5 92.916\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "with torch.no_grad():\n",
    "    evaluate(model, criterion, data_loader_test, device=\"cuda\", print_freq=20)\n",
    "    \n",
    "# Save the model\n",
    "torch.save(model.state_dict(), \"/tmp/quant_resnet50-calibrated.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We can also try different calibrations and see which one works the best"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test:  [ 0/98]  eta: 0:05:27  loss: 0.6037 (0.6037)  acc1: 84.9609 (84.9609)  acc5: 95.3125 (95.3125)  time: 3.3411  data: 2.6190  max mem: 5882\n",
      "Test:  [20/98]  eta: 0:01:06  loss: 0.6760 (0.7041)  acc1: 81.2500 (81.7522)  acc5: 95.7031 (95.4892)  time: 0.7243  data: 0.0002  max mem: 5882\n",
      "Test:  [40/98]  eta: 0:00:45  loss: 0.7241 (0.7351)  acc1: 79.1016 (80.7784)  acc5: 95.8984 (95.4459)  time: 0.7243  data: 0.0002  max mem: 5882\n",
      "Test:  [60/98]  eta: 0:00:29  loss: 1.1162 (0.8793)  acc1: 71.4844 (77.6383)  acc5: 90.8203 (93.7948)  time: 0.7204  data: 0.0002  max mem: 5882\n",
      "Test:  [80/98]  eta: 0:00:13  loss: 1.1498 (0.9603)  acc1: 71.4844 (76.0368)  acc5: 89.4531 (92.7156)  time: 0.7164  data: 0.0002  max mem: 5882\n",
      "Test: Total time: 0:01:12\n",
      " * Acc@1 75.438 Acc@5 92.486\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    compute_amax(model, method=\"percentile\", percentile=99.9)\n",
    "    evaluate(model, criterion, data_loader_test, device=\"cuda\", print_freq=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse calibration\n",
      "Test:  [ 0/98]  eta: 0:06:34  loss: 0.5700 (0.5700)  acc1: 85.1562 (85.1562)  acc5: 96.2891 (96.2891)  time: 4.0243  data: 3.3231  max mem: 5882\n",
      "Test:  [20/98]  eta: 0:01:08  loss: 0.6758 (0.6838)  acc1: 82.8125 (82.5707)  acc5: 96.0938 (95.7868)  time: 0.7204  data: 0.0002  max mem: 5882\n",
      "Test:  [40/98]  eta: 0:00:46  loss: 0.7047 (0.7163)  acc1: 80.2734 (81.4834)  acc5: 96.2891 (95.7746)  time: 0.7178  data: 0.0002  max mem: 5882\n",
      "Test:  [60/98]  eta: 0:00:29  loss: 1.1127 (0.8585)  acc1: 71.0938 (78.3395)  acc5: 90.8203 (94.1278)  time: 0.7192  data: 0.0002  max mem: 5882\n",
      "Test:  [80/98]  eta: 0:00:13  loss: 1.1261 (0.9367)  acc1: 72.6562 (76.7530)  acc5: 89.8438 (93.1785)  time: 0.7176  data: 0.0002  max mem: 5882\n",
      "Test: Total time: 0:01:13\n",
      " * Acc@1 76.186 Acc@5 92.926\n",
      "entropy calibration\n",
      "Test:  [ 0/98]  eta: 0:05:28  loss: 0.5648 (0.5648)  acc1: 85.3516 (85.3516)  acc5: 96.0938 (96.0938)  time: 3.3558  data: 2.6268  max mem: 5882\n",
      "Test:  [20/98]  eta: 0:01:05  loss: 0.6724 (0.6815)  acc1: 82.8125 (82.5428)  acc5: 95.8984 (95.7589)  time: 0.7196  data: 0.0002  max mem: 5882\n",
      "Test:  [40/98]  eta: 0:00:45  loss: 0.7090 (0.7149)  acc1: 80.6641 (81.4929)  acc5: 96.0938 (95.7269)  time: 0.7214  data: 0.0002  max mem: 5882\n",
      "Test:  [60/98]  eta: 0:00:29  loss: 1.1077 (0.8571)  acc1: 72.0703 (78.3779)  acc5: 90.6250 (94.0798)  time: 0.7198  data: 0.0002  max mem: 5882\n",
      "Test:  [80/98]  eta: 0:00:13  loss: 1.1253 (0.9356)  acc1: 72.2656 (76.7626)  acc5: 90.0391 (93.1231)  time: 0.7192  data: 0.0002  max mem: 5882\n",
      "Test: Total time: 0:01:12\n",
      " * Acc@1 76.206 Acc@5 92.900\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    for method in [\"mse\", \"entropy\"]:\n",
    "        print(F\"{method} calibration\")\n",
    "        compute_amax(model, method=method)\n",
    "        evaluate(model, criterion, data_loader_test, device=\"cuda\", print_freq=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
