notebooks/114-quantization-simplified-mode/114-quantization-simplified-mode.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# INT8 Quantization with Post-training Optimization Tool (POT) in Simplified Mode tutorial\n",
    "\n",
    "This tutorial shows how to quantize a [ResNet20](https://github.com/chenyaofo/pytorch-cifar-models) image classification model, trained on [CIFAR10 ](http://pytorch.org/vision/main/generated/torchvision.datasets.CIFAR10.html) dataset, using the Post-Training Optimization Tool (POT) in Simplified Mode.\n",
    "\n",
    "Simplified Mode is designed to make the data preparation step easier, before model optimization. The mode is represented by an implementation of the engine interface in the POT API in OpenVINO™. It enables reading data from an arbitrary folder specified by the user. Currently, Simplified Mode is available only for image data in PNG or JPEG formats, stored in a single folder.\n",
    "\n",
    "> **Note:** This mode cannot be used with the accuracy-aware method. It is not possible to control accuracy after optimization using this mode. However, Simplified Mode can be useful for estimating performance improvements when optimizing models.\n",
    "\n",
    "This tutorial includes the following steps:\n",
    "\n",
    "- Downloading and saving the CIFAR10 dataset.\n",
    "- Preparing the model for quantization.\n",
    "- Compressing the prepared model.\n",
    "- Measuring and comparing the performance of the original and quantized models.\n",
    "- Demonstrating the use of the quantized model for image classification.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import warnings\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms as T\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from openvino.runtime import Core, Tensor\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Set the data and model directories\n",
    "MODEL_DIR = 'model'\n",
    "CALIB_DIR = 'calib'\n",
    "CIFAR_DIR = '../data/datasets/cifar10'\n",
    "CALIB_SET_SIZE = 300\n",
    "MODEL_NAME = 'resnet20'\n",
    "\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)\n",
    "os.makedirs(CALIB_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the calibration dataset\n",
    "The following steps are required to prepare the calibration dataset:\n",
    "- Download the CIFAR10 dataset from [Torchvision.datasets repository](https://pytorch.org/vision/stable/datasets.html).\n",
    "- Save the selected number of elements from this dataset as `.png` images in a separate folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = T.Compose([T.ToTensor()])\n",
    "dataset = CIFAR10(root=CIFAR_DIR, train=False, transform=transform, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pil_converter = T.ToPILImage(mode=\"RGB\")\n",
    "\n",
    "for idx, info in enumerate(dataset):\n",
    "    im = info[0]\n",
    "    if idx >= CALIB_SET_SIZE:\n",
    "        break\n",
    "    label = info[1]\n",
    "    pil_converter(im.squeeze(0)).save(Path(CALIB_DIR) / f'{label}_{idx}.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the Model\n",
    "\n",
    "Model preparation includes the following steps:\n",
    "- Download PyTorch model from Torchvision repository.\n",
    "- Convert the model to ONNX format.\n",
    "- Run Model Optimizer to convert ONNX to OpenVINO Intermediate Representation (OpenVINO IR)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.hub.load(\"chenyaofo/pytorch-cifar-models\", \"cifar10_resnet20\", pretrained=True, skip_validation=True)\n",
    "dummy_input = torch.randn(1, 3, 32, 32)\n",
    "\n",
    "onnx_model_path = Path(MODEL_DIR) / '{}.onnx'.format(MODEL_NAME)\n",
    "ir_model_xml = onnx_model_path.with_suffix('.xml')\n",
    "ir_model_bin = onnx_model_path.with_suffix('.bin')\n",
    "\n",
    "torch.onnx.export(model, dummy_input, onnx_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, convert this model into the OpenVINO IR using Model Optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo -m $onnx_model_path  --output_dir $MODEL_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compression stage\n",
    "Compress the model with the following command:\n",
    "  \n",
    "`pot -q default -m <path_to_xml> -w <path_to_bin> --engine simplified --data-source <path_to_data>`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pot -q default -m $ir_model_xml -w $ir_model_bin --engine simplified --data-source $CALIB_DIR --output-dir compressed --direct-dump --name $MODEL_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Performance of the Original and Quantized Models\n",
    "\n",
    "Finally, measure the inference performance of the `FP32` and `INT8` models, using [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) - an inference performance measurement tool in OpenVINO.\n",
    "\n",
    "> **Note**: For more accurate performance, it is recommended to run benchmark_app in a terminal/command prompt after closing other applications. Run `benchmark_app -m model.xml -d CPU` to benchmark async inference on CPU for one minute. Change CPU to GPU to benchmark on GPU. Run `benchmark_app --help` to see an overview of all command-line options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model_path = Path('compressed/optimized')\n",
    "optimized_model_xml = optimized_model_path / '{}.xml'.format(MODEL_NAME)\n",
    "optimized_model_bin = optimized_model_path / '{}.bin'.format(MODEL_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference FP32 model (OpenVINO IR)\n",
    "!benchmark_app -m $ir_model_xml -d CPU -api async"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference INT8 model (OpenVINO IR)\n",
    "!benchmark_app -m $optimized_model_xml -d CPU -api async"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demonstration of the results\n",
    "\n",
    "This section demonstrates how to use the compressed model by running the optimized model on a subset of images from the CIFAR10 dataset and shows predictions, using the model.\n",
    "\n",
    "The first step is to load the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "\n",
    "compiled_model = ie.compile_model(str(optimized_model_xml))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define all possible labels from the CIFAR10 dataset.\n",
    "labels_names = [\"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\", \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"]\n",
    "all_images = []\n",
    "all_labels = []\n",
    "\n",
    "# Get all images and their labels. \n",
    "for batch in dataset:\n",
    "    all_images.append(torch.unsqueeze(batch[0], 0))\n",
    "    all_labels.append(batch[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below defines the function that shows the images and their labels, using the indexes and two lists created in the previous step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_pictures(indexes: list, images=all_images, labels=all_labels):\n",
    "    \"\"\"Plot images with the specified indexes.\n",
    "    :param indexes: a list of indexes of images to be displayed.\n",
    "    :param images: a list of images from the dataset.\n",
    "    :param labels: a list of labels for each image.\n",
    "    \"\"\"\n",
    "    num_pics = len(indexes)\n",
    "    _, axarr = plt.subplots(1, num_pics)\n",
    "    for idx, im_idx in enumerate(indexes):\n",
    "        assert idx < 10000, 'Cannot get such index, there are only 10000'\n",
    "        pic = np.rollaxis(images[im_idx].squeeze().numpy(), 0, 3)\n",
    "        axarr[idx].imshow(pic)\n",
    "        axarr[idx].set_title(labels_names[labels[im_idx]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the code below, to define a function that uses the optimized model to obtain predictions for the selected images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_on_images(net, indexes: list, images=all_images):\n",
    "    \"\"\" Inference model on a set of images.\n",
    "    :param net: model on which do inference\n",
    "    :param indexes: a list of indexes of images to infer on.\n",
    "    :param images: a list of images from the dataset.\n",
    "    \"\"\"\n",
    "    predicted_labels = []\n",
    "    infer_request = net.create_infer_request()\n",
    "    for idx in indexes:\n",
    "        assert idx < 10000, 'Cannot get such index, there are only 10000'\n",
    "        input_tensor = Tensor(array=images[idx].detach().numpy(), shared_memory=True)\n",
    "        infer_request.set_input_tensor(input_tensor)\n",
    "        infer_request.start_async()\n",
    "        infer_request.wait()\n",
    "        output = infer_request.get_output_tensor()\n",
    "        result = list(output.data)\n",
    "        result = labels_names[np.argmax(result[0])]\n",
    "        predicted_labels.append(result)\n",
    "    return predicted_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexes_to_infer = [0, 1, 2]  # to plot specify indexes\n",
    "\n",
    "plot_pictures(indexes_to_infer)\n",
    "\n",
    "results_quanized = infer_on_images(compiled_model, indexes_to_infer)\n",
    "\n",
    "print(f\"Image labels using the quantized model : {results_quanized}.\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "K5HPrY_d-7cV",
    "E01dMaR2_AFL",
    "qMnYsGo9_MA8",
    "L0tH9KdwtHhV"
   ],
   "name": "INT8 Quantization with POT in Simplified Mode tutorial",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}