optimization/accuracyQuantization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Google Summer of Code 2022: Train a DL model for synthetic data generation for model optimization\n",
    "\n",
    "This notebook performs the [accuracy-control quantization](https://docs.openvino.ai/latest/pot_accuracyaware_usage.html) method of OpenVINO's Post-training Optimization Tool on the two models, namely ShuffleNetV2 and RepVGG, that showcased accuracy degradation with the default quantization method, as part of the GSoC22 project \"Train a DL model for synthetic data generation for model optimization\", which is implemented under the auspices of Intel's OpenVINO Toolkit organization. The performance of the models is then evaluated on the classification task on the CIFAR-10 test set. The selected CV models are pre-trained on CIFAR-10 and obtained from the [chenyaofo/pytorch-cifar-models](https://github.com/chenyaofo/pytorch-cifar-models) repository on GitHub."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import ipywidgets as widgets\n",
    "import numpy as np\n",
    "import PIL\n",
    "import torch\n",
    "from IPython.display import Markdown, display\n",
    "from matplotlib import pyplot as plt\n",
    "from openvino.runtime import Core\n",
    "from openvino.tools.pot import (IEEngine, compress_model_weights,\n",
    "                                create_pipeline, load_model, save_model)\n",
    "from openvino.tools.pot.api import DataLoader, Metric\n",
    "from sklearn.metrics import accuracy_score\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import CIFAR10, ImageFolder\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data\n",
    "Load the four different datasets:\n",
    "* Official CIFAR-10 training set\n",
    "* FakeCIFAR generated by the StyleGAN2-ADA model\n",
    "* FakeCIFAR generated by the DiStyleGAN model\n",
    "* Fractal Images generated using the model from [Datumaro's repository](https://github.com/openvinotoolkit/datumaro) on GitHub\n",
    "\n",
    "and create a DataLoader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CIFARLoader(DataLoader):\n",
    "    \"\"\"\n",
    "    DataLoader for image data that is stored in a directory per category.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data_source) -> None:\n",
    "        \"\"\"\n",
    "        - data_source: dataset for which to create loader\n",
    "        \"\"\"\n",
    "        self.dataset = data_source\n",
    "        self.class_names = self.dataset.classes\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"\n",
    "        Returns the number of elements in the dataset\n",
    "        \"\"\"\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def __getitem__(self, index: int) -> tuple:\n",
    "        \"\"\"\n",
    "        Get item from self.dataset at the specified index.\n",
    "        Returns (annotation, image), where annotation is a tuple (index, class_index)\n",
    "        and image a preprocessed image in network shape\n",
    "        \"\"\"\n",
    "        if index >= len(self):\n",
    "            raise IndexError\n",
    "        image, annotation = self.dataset[index]\n",
    "        return (index, annotation), torch.unsqueeze(image, dim=0).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(),\n",
    "     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CIFAR10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets['CIFAR10'] = CIFAR10(\".\",\n",
    "    download=True, transform=transform)\n",
    "\n",
    "cifar10_ds_test =  CIFAR10(\".\", train=False,\n",
    "    download=True, transform=transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FakeCIFAR StyleGAN2-ADA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets['StyleGAN2-ADA'] = ImageFolder(\"datasets/StyleGAN2-ADA/\", transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FakeCIFAR DistyleGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets['DiStyleGAN'] = ImageFolder(\"datasets/DiStyleGAN/\", transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fractal Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets['Fractal'] = ImageFolder(\"datasets/Fractal/\", transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the PyTorch model\n",
    "Select the PyTorch model to be 8-bit quantized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selection = widgets.Dropdown(\n",
    "    options=['ShuffleNetV2', 'RepVGG'],\n",
    "    value='ShuffleNetV2',\n",
    "    description='Model:',\n",
    "    disabled=False,\n",
    ")\n",
    "display(selection)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the model from torch.hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_names = {\n",
    "    'ShuffleNetV2' : 'shufflenetv2_x2_0',\n",
    "    'RepVGG' : 'repvgg_a2'\n",
    "} \n",
    "\n",
    "model_name = model_names[selection.value]\n",
    "model = torch.hub.load(\"chenyaofo/pytorch-cifar-models\", f\"cifar10_{model_name}\", pretrained=True)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the PyTorch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.perf_counter()\n",
    "predictions = []\n",
    "true_labels = []\n",
    "with torch.no_grad():\n",
    "    for img, _class in cifar10_ds_test:\n",
    "        img = img.unsqueeze(dim=0)\n",
    "        res = model(img)\n",
    "        predictions.append(int(np.argmax(res)))\n",
    "        true_labels.append(_class)\n",
    "\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(\n",
    "    f\"Optimized model in Inference Engine/CPU: {time_ir/len(cifar10_ds_test):.3f} \"\n",
    "    f\"seconds per image, FPS: {len(cifar10_ds_test)/time_ir:.2f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_score(true_labels, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "directory = Path(\"optimizations\")\n",
    "directory.mkdir(exist_ok=True)\n",
    "\n",
    "save = directory / Path(f\"{selection.value}\")\n",
    "\n",
    "# Paths where PyTorch, ONNX and OpenVINO IR models will be stored\n",
    "onnx_path = save.with_suffix(\".onnx\")\n",
    "ir_path = save.with_suffix(\".xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ONNX conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not onnx_path.exists():\n",
    "    dummy_input = torch.randn(1, 3, 32, 32)\n",
    "    \n",
    "    torch.onnx.export(\n",
    "        model,\n",
    "        (dummy_input, ),\n",
    "        onnx_path,\n",
    "        opset_version=11\n",
    "    )\n",
    "    print(f\"ONNX model exported to {onnx_path}.\")\n",
    "else:\n",
    "    print(f\"ONNX model {onnx_path} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert ONNX to OpenVINO IR format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct the command for Model Optimizer\n",
    "mo_command = f\"\"\"mo\n",
    "                 --input_model \"{onnx_path}\"\n",
    "                 --output_dir \"{save.parent}\"\n",
    "                 \"\"\"\n",
    "mo_command = \" \".join(mo_command.split())\n",
    "print(\"Model Optimizer command to convert the ONNX model to OpenVINO:\")\n",
    "display(Markdown(f\"`{mo_command}`\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not ir_path.exists():\n",
    "    print(\"Exporting ONNX model to IR... This may take a few minutes.\")\n",
    "    mo_result = %sx $mo_command\n",
    "    print(\"\\n\".join(mo_result))\n",
    "else:\n",
    "    print(f\"IR model {ir_path} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate IR model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the optimized model and get the names of the input and output layer\n",
    "ie = Core()\n",
    "model_pot = ie.read_model(model=f\"optimizations/{selection.value}.xml\")\n",
    "compiled_model_pot = ie.compile_model(model=model_pot, device_name=\"CPU\")\n",
    "input_layer = compiled_model_pot.input(0)\n",
    "output_layer = compiled_model_pot.output(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.perf_counter()\n",
    "predictions = []\n",
    "true_labels = []\n",
    "for img, _class in cifar10_ds_test:\n",
    "    img = img.unsqueeze(dim=0)\n",
    "    res = compiled_model_pot([img])[output_layer]\n",
    "    predictions.append(int(np.argmax(res)))\n",
    "    true_labels.append(_class)\n",
    "\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(\n",
    "    f\"Optimized model in Inference Engine/CPU: {time_ir/len(cifar10_ds_test):.3f} \"\n",
    "    f\"seconds per image, FPS: {len(cifar10_ds_test)/time_ir:.2f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_score(true_labels, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define accuracy metric\n",
    "More information about the accuracy metric can be found [here](https://docs.openvino.ai/latest/pot_accuracyaware_usage.html#define-accuracy-metric)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Accuracy(Metric):\n",
    "\n",
    "    # Required methods\n",
    "    def __init__(self, top_k=1):\n",
    "        super().__init__()\n",
    "        self._top_k = top_k\n",
    "        self._name = 'accuracy@top{}'.format(self._top_k)\n",
    "        self._matches = [] # container of the results\n",
    "\n",
    "    @property\n",
    "    def value(self):\n",
    "        \"\"\" Returns accuracy metric value for all model outputs. \"\"\"\n",
    "        return {self._name: self._matches[-1]}\n",
    "\n",
    "    @property\n",
    "    def avg_value(self):\n",
    "        \"\"\" Returns accuracy metric value for all model outputs. \"\"\"\n",
    "        return {self._name: np.ravel(self._matches).mean()}\n",
    "\n",
    "    def update(self, output, target):\n",
    "        \"\"\" Updates prediction matches.\n",
    "        :param output: model output\n",
    "        :param target: annotations\n",
    "        \"\"\"\n",
    "        if len(output) > 1:\n",
    "            raise Exception('The accuracy metric cannot be calculated '\n",
    "                            'for a model with multiple outputs')\n",
    "        if isinstance(target, dict):\n",
    "            target = list(target.values())\n",
    "        predictions = np.argsort(output[0], axis=1)[:, -self._top_k:]\n",
    "        match = [float(t in predictions[i]) for i, t in enumerate(target)]\n",
    "\n",
    "        self._matches.append(match)\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\" Resets collected matches \"\"\"\n",
    "        self._matches = []\n",
    "\n",
    "    def get_attributes(self):\n",
    "        \"\"\"\n",
    "        Returns a dictionary of metric attributes {metric_name: {attribute_name: value}}.\n",
    "        Required attributes: 'direction': 'higher-better' or 'higher-worse'\n",
    "                             'type': metric type\n",
    "        \"\"\"\n",
    "        return {self._name: {'direction': 'higher-better',\n",
    "                             'type': 'accuracy'}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Select calibration dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calibration = widgets.Dropdown(\n",
    "    options=['CIFAR10', 'StyleGAN2-ADA', 'DiStyleGAN', 'Fractal'],\n",
    "    value='CIFAR10',\n",
    "    description='Model:',\n",
    "    disabled=False,\n",
    ")\n",
    "display(calibration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_bin = ir_path.with_suffix('.bin')\n",
    "\n",
    "# Model config specifies the model name and paths to model .xml and .bin file\n",
    "model_config = {\n",
    "    \"model_name\": \"model\",\n",
    "    \"model\": ir_path,\n",
    "    \"weights\": path_to_bin,\n",
    "}\n",
    "\n",
    "# Engine config\n",
    "engine_config = {\"device\": \"CPU\"}\n",
    "\n",
    "algorithms = [\n",
    "    {\n",
    "        \"name\": \"AccuracyAwareQuantization\",\n",
    "        \"params\": {\n",
    "            \"target_device\": \"ANY\",\n",
    "            'maximal_drop': 0.02\n",
    "        },\n",
    "    }\n",
    "]\n",
    "\n",
    "# Step 1: Implement and create user's data loader\n",
    "data_loader = CIFARLoader(datasets[calibration.value])\n",
    "\n",
    "# Step 2: Metric\n",
    "metric = Accuracy()\n",
    "\n",
    "# Step 3: Load model\n",
    "model = load_model(model_config=model_config)\n",
    "\n",
    "# Step 4: Initialize the engine for metric calculation and statistics collection.\n",
    "engine = IEEngine(config=engine_config, data_loader=data_loader, metric=metric)\n",
    "\n",
    "# Step 5: Create a pipeline of compression algorithms and run it.\n",
    "pipeline = create_pipeline(algorithms, engine)\n",
    "compressed_model = pipeline.run(model=model)\n",
    "\n",
    "# Step 6 (Optional): Compress model weights to quantized precision\n",
    "#                     to reduce the size of the final .bin file.\n",
    "compress_model_weights(compressed_model)\n",
    "\n",
    "# Step 7: Save the compressed model to the desired path.\n",
    "# Set save_path to the directory where the model should be saved\n",
    "compressed_model_paths = save_model(\n",
    "    model=compressed_model,\n",
    "    save_path=save,\n",
    "    model_name=f\"optimized_model_{calibration.value}_acc\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate quantized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the optimized model and get the names of the input and output layer\n",
    "ie = Core()\n",
    "model_pot = ie.read_model(model=f\"optimizations/{selection.value}/optimized_model_{calibration.value}_acc.xml\")\n",
    "compiled_model_pot = ie.compile_model(model=model_pot, device_name=\"CPU\")\n",
    "input_layer = compiled_model_pot.input(0)\n",
    "output_layer = compiled_model_pot.output(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.perf_counter()\n",
    "predictions = []\n",
    "true_labels = []\n",
    "for img, _class in cifar10_ds_test:\n",
    "    img = img.unsqueeze(dim=0)\n",
    "    res = compiled_model_pot([img])[output_layer]\n",
    "    predictions.append(int(np.argmax(res)))\n",
    "    true_labels.append(_class)\n",
    "\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(\n",
    "    f\"Optimized model in Inference Engine/CPU: {time_ir/len(cifar10_ds_test):.3f} \"\n",
    "    f\"seconds per image, FPS: {len(cifar10_ds_test)/time_ir:.2f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_score(true_labels, predictions)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openvino",
   "language": "python",
   "name": "openvino"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}