1. Tutorial - automatic quantization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatic quantization\n",
    "\n",
    "This notebook demonstrates simple end-two-end pipeline for MobileNetV2 quantization.\n",
    "\n",
    "Our quantization process consists of quantized model calibration, quantization threshold adjustment and weight fine-tuning using distillation. Finally, we demonstrate inference of our quantized model using ONNX Runtime framework.\n",
    "\n",
    "### Main chapters of this notebook:\n",
    "1. Setup the environment\n",
    "1. Prepare dataset and create dataloaders\n",
    "1. Evaluate pretrained MobileNetV2 from torchvision\n",
    "1. End2end quantization with our framework\n",
    "1. Inference using ONNX Runtime with TensorRT Execution Provider\n",
    "\n",
    "Before running this example make sure that TensorRT supports your GPU for int8 inference  (``cuda compute capability`` > 6.1, as described [here](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#hardware-precision-matrix))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the environment\n",
    "\n",
    "First, let's set up the environment and make some common imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'\n",
    "# You may need to uncomment and change this variable to match free GPU index\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Common:\n",
    "import itertools\n",
    "import numpy as np\n",
    "import torch\n",
    "from pathlib import Path\n",
    "from torch import nn\n",
    "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
    "from torch.optim import RAdam\n",
    "from tqdm.auto import tqdm\n",
    "from tutorial_utils.dataset import create_imagenet10k_dataloaders\n",
    "from tutorial_utils.train import accuracy\n",
    "\n",
    "# Quantization:\n",
    "from enot.quantization import RMSELoss\n",
    "from enot.quantization import TensorRTFakeQuantizedModel\n",
    "from enot.quantization import calibrate\n",
    "from enot.quantization import distill\n",
    "\n",
    "# ONNX Runtime inference:\n",
    "from tutorial_utils.inference import create_onnxruntime_session\n",
    "import onnxsim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define model evaluation function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function can evaluate both nn.Modules and executable functions.\n",
    "def eval_model(model_fn, dataloader):\n",
    "    if isinstance(model_fn, nn.Module):\n",
    "        model_fn.eval()\n",
    "\n",
    "    total = 0\n",
    "    total_loss = 0.0\n",
    "    total_correct = 0.0\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in tqdm(dataloader):\n",
    "            n = inputs.shape[0]\n",
    "\n",
    "            pred_labels = model_fn(inputs)\n",
    "            batch_loss = criterion(pred_labels, labels)\n",
    "            batch_accuracy = accuracy(pred_labels, labels)\n",
    "\n",
    "            total += n\n",
    "            total_loss += batch_loss.item() * n\n",
    "            total_correct += batch_accuracy.item() * n\n",
    "\n",
    "    return total_loss / total, total_correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### In the following cell we setup all necessary dirs\n",
    "\n",
    "* `HOME_DIR` - experiments home directory\n",
    "* `DATASETS_DIR` - root directory for datasets (imagenette2, ...)\n",
    "* `PROJECT_DIR` - project directory to save training logs, checkpoints, ...\n",
    "* `ONNX_MODEL_PATH` - onnx model path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME_DIR = Path.home() / '.optimization_experiments'\n",
    "DATASETS_DIR = HOME_DIR / 'datasets'\n",
    "PROJECT_DIR = HOME_DIR / 'enot-lite_quantization'\n",
    "ONNX_MODEL_PATH = PROJECT_DIR / 'mobilenetv2.onnx'\n",
    "\n",
    "HOME_DIR.mkdir(exist_ok=True)\n",
    "DATASETS_DIR.mkdir(exist_ok=True)\n",
    "PROJECT_DIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare dataset and create dataloaders\n",
    "\n",
    "We will use Imagenet-10k dataset in this example.\n",
    "\n",
    "Imagenet-10k dataset is a subsample of [Imagenet](https://image-net.org/challenges/LSVRC/index.php) dataset. It contains 5000 training images and 5000 validation images. Training images are uniformly gathered from the original training set, and validation images are gathered from the original validation set, 5 per each class.\n",
    "\n",
    "`create_imagenet10k_dataloaders` function prepares datasets for you in this example; specifically, it:\n",
    "1. downloads and unpacks dataset into `DATASETS_DIR`;\n",
    "1. creates and returns train and validation dataloaders.\n",
    "\n",
    "The two parts of the dataset:\n",
    "* train: for quantization procedure (`DATASETS_DIR`/imagenet10k/train/)\n",
    "* validation: for model validation (`DATASETS_DIR`/imagenet10k/val/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader, validation_dataloader = create_imagenet10k_dataloaders(\n",
    "    dataset_root_dir=DATASETS_DIR,\n",
    "    input_size=224,\n",
    "    batch_size=25,\n",
    "    num_workers=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate pretrained MobileNetV2 from torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models.mobilenetv2 import mobilenet_v2\n",
    "\n",
    "regular_model = mobilenet_v2(pretrained=True).cuda()\n",
    "\n",
    "# Turning off FullyConnected layer dropout.\n",
    "# This is required to stabilize fine-tuning procedure.\n",
    "regular_model.classifier[0].p = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loss, val_accuracy = eval_model(regular_model, validation_dataloader)\n",
    "print(f'Regular (non-quantized) model: accuracy={val_accuracy:.3f}, loss={val_loss:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End2end quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- wrap `regular_model` to `TensorRTFakeQuantizedModel`\n",
    "- calibrate quantization threshold using `calibration` context\n",
    "- distill quantization threshold and scale-factors using `distill` context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fake_quantized_model = TensorRTFakeQuantizedModel(regular_model).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calibrate quantization thresholds using 10 batches.\n",
    "with torch.no_grad(), calibrate(fake_quantized_model):\n",
    "    for batch in itertools.islice(train_dataloader, 10):\n",
    "        batch = batch[0].cuda()\n",
    "        fake_quantized_model(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Distill quantization thresholds and scale-factors using RMSE loss for 5 epochs.\n",
    "n_epochs = 5\n",
    "\n",
    "with distill(fq_model=fake_quantized_model, tune_weight_scale_factors=True) as (qdistill_model, params):\n",
    "    optimizer = RAdam(params=params, lr=0.005, betas=(0.9, 0.95))\n",
    "    scheduler = CosineAnnealingLR(optimizer=optimizer, T_max=len(train_dataloader) * n_epochs)\n",
    "    distillation_criterion = RMSELoss()\n",
    "\n",
    "    for _ in range(n_epochs):\n",
    "        for batch in (tqdm_it := tqdm(train_dataloader)):\n",
    "            batch = batch[0].cuda()\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss: torch.Tensor = torch.tensor(0.0).cuda()\n",
    "            for student_output, teacher_output in qdistill_model(batch):\n",
    "                loss += distillation_criterion(student_output, teacher_output)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "\n",
    "            tqdm_it.set_description(f'loss: {loss.item():.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fake_quantized_model.enable_quantization_mode(True)\n",
    "val_loss, val_accuracy = eval_model(fake_quantized_model, validation_dataloader)\n",
    "print(f'Optimized quantized model: accuracy={val_accuracy:.3f}, loss={val_loss:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference using ONNX Runtime with TensorRT Execution Provider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(\n",
    "    model=fake_quantized_model.cpu(),\n",
    "    args=torch.zeros(25, 3, 224, 224),\n",
    "    f='exported_model.onnx',\n",
    "    opset_version=13,\n",
    "    input_names=['input'],\n",
    "    output_names=['output'],\n",
    ")\n",
    "\n",
    "proto, _ = onnxsim.simplify('exported_model.onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize ONNX Runtime inference session with TensorRT Execution Provider:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()  # Empty PyTorch CUDA cache before running ONNX Runtime.\n",
    "\n",
    "sess = create_onnxruntime_session(\n",
    "    proto=proto,\n",
    "    input_sample=torch.zeros(25, 3, 224, 224, device='cuda'),\n",
    "    output_shape=(25, 1000),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate quantized model on TensorRT Execution Provider:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(inputs):\n",
    "    return sess(inputs)\n",
    "\n",
    "\n",
    "val_loss, val_accuracy = eval_model(model_fn, validation_dataloader)\n",
    "print(f'Quantized model with fine-tuned weights with TRT: accuracy={val_accuracy:.3f}, loss={val_loss:.3f}')"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 3000
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}