notebooks/113-image-classification-quantization/113-image-classification-quantization.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rQc-wXjqrEuR"
      },
      "source": [
        "# Quantization of Image Classification Models\n",
        "\n",
        "This tutorial demostrates how to apply INT8 quantization to Image Classification model using [Post-training Optimization Tool API](../../compression/api/README.md). The Mobilenet V2 model trained on Cifar10 dataset is used as an example. The code of this tutorial is designed to be extandable to custom model and dataset. It is assumed that OpenVINO is already installed. This tutorial consists of the following steps:\n",
        "- Prepare the model for quantization\n",
        "- Define data loading and accuracy validation functionality\n",
        "- Run optimization pipeline\n",
        "- Compare accuracy of the original and quantized models\n",
        "- Compare performance of the original and quantized models\n",
        "- Compare results on one picture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4cSNQWdbSyeo"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pathlib import Path\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch\n",
        "from addict import Dict\n",
        "from compression.api import DataLoader, Metric\n",
        "from compression.engines.ie_engine import IEEngine\n",
        "from compression.graph import load_model, save_model\n",
        "from compression.graph.model_utils import compress_model_weights\n",
        "from compression.pipeline.initializer import create_pipeline\n",
        "from openvino.runtime import Core\n",
        "from torchvision import transforms\n",
        "from torchvision.datasets import CIFAR10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Set the data and model directories\n",
        "DATA_DIR = 'data'\n",
        "MODEL_DIR = 'model'\n",
        "\n",
        "os.makedirs(DATA_DIR, exist_ok=True)\n",
        "os.makedirs(MODEL_DIR, exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D-frbVLKrkmv"
      },
      "source": [
        "## Prepare the Model\n",
        "Model preparation stage has the following steps:\n",
        "- Download PyTorch model from Torchvision repository\n",
        "- Convert it to ONNX format\n",
        "- Run OpenVINO Model Optimizer tool to convert ONNX to OpenVINO Intermediate Representation (IR)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f7i6dWUmhloy"
      },
      "outputs": [],
      "source": [
        "model = torch.hub.load(\"chenyaofo/pytorch-cifar-models\", \"cifar10_mobilenetv2_x1_0\", pretrained=True)\n",
        "model.eval()\n",
        "\n",
        "dummy_input = torch.randn(1, 3, 32, 32)\n",
        "\n",
        "onnx_model_path = Path(MODEL_DIR) / 'mobilenet_v2.onnx'\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, verbose=True)\n",
        "\n",
        "# Run OpenVINO Model Optimization tool to convert ONNX to OpenVINO IR\n",
        "!mo --framework=onnx --data_type=FP16 --input_shape=[1,3,32,32] -m $onnx_model_path  --output_dir $MODEL_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ynLvh8rNc2wv"
      },
      "source": [
        "## Define Data Loader\n",
        "At this step the `DataLoader` interface from POT API is implemented."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])\n",
        "dataset = CIFAR10(root=DATA_DIR, train=False, transform=transform, download=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DErQofk8tO6c",
        "test_replace": {
          "len(self.labels)": "100"
        }
      },
      "outputs": [],
      "source": [
        "# create DataLoader from CIFAR10 dataset\n",
        "class CifarDataLoader(DataLoader):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        \"\"\"\n",
        "        Initialize config and dataset.\n",
        "        :param config: created config with DATA_DIR path.\n",
        "        \"\"\"\n",
        "        if not isinstance(config, Dict):\n",
        "            config = Dict(config)\n",
        "        super().__init__(config)\n",
        "        self.indexes, self.pictures, self.labels = self.load_data(dataset)\n",
        "        \n",
        "    def __len__(self):\n",
        "        return len(self.labels)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        \"\"\"\n",
        "        Return one sample of index, label and picture.\n",
        "        :param index: index of the taken sample.\n",
        "        \"\"\"\n",
        "        if index >= len(self):\n",
        "            raise IndexError\n",
        "\n",
        "        return (self.indexes[index], self.labels[index]), self.pictures[index].numpy()\n",
        "\n",
        "    def load_data(self, dataset):\n",
        "        \"\"\"\n",
        "        Load dataset in needed format. \n",
        "        :param dataset:  downloaded dataset.\n",
        "        \"\"\"\n",
        "        pictures, labels, indexes = [], [], []\n",
        "        \n",
        "        for idx, sample in enumerate(dataset):\n",
        "            pictures.append(sample[0])\n",
        "            labels.append(sample[1])\n",
        "            indexes.append(idx)\n",
        "\n",
        "        return indexes, pictures, labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Re9-YhbBddh3"
      },
      "source": [
        "## Define Accuracy Metric Calculation\n",
        "At this step the `Metric` interface for accuracy Top-1 metric is implemented. It is used for validating accuracy of quantized model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GB8L492ztZEC"
      },
      "outputs": [],
      "source": [
        "# Custom implementation of classification accuracy metric.\n",
        "\n",
        "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 = []\n",
        "\n",
        "    @property\n",
        "    def value(self):\n",
        "        \"\"\" Returns accuracy metric value for the last model output. \"\"\"\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": {
        "id": "CclWk-fVd9Wi"
      },
      "source": [
        "## Run Quantization Pipeline and compare the accuracy of the original and quantized models\n",
        "Here we define a configuration for our quantization pipeline and run it. \n",
        "\n",
        "NOTE: we use built-in `IEEngine` implementation of the `Engine` interface from the POT API for model inference. `IEEngine` is built on top of OpenVINO Python* API for inference and provides basic functionality for inference of simple models. If you have a more complicated inference flow for your model/models you should create your own implementation of `Engine` interface, for example by inheriting from `IEEngine` and extending it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PiAvrwo0tr6Z"
      },
      "outputs": [],
      "source": [
        "model_config = Dict({\n",
        "    'model_name': 'mobilenet_v2',\n",
        "    'model': ir_model_xml,\n",
        "    'weights': ir_model_bin\n",
        "})\n",
        "engine_config = Dict({\n",
        "    'device': 'CPU',\n",
        "    'stat_requests_number': 2,\n",
        "    'eval_requests_number': 2\n",
        "})\n",
        "dataset_config = {\n",
        "    'data_source': DATA_DIR\n",
        "}\n",
        "algorithms = [\n",
        "    {\n",
        "        'name': 'DefaultQuantization',\n",
        "        'params': {\n",
        "            'target_device': 'CPU',\n",
        "            'preset': 'performance',\n",
        "            'stat_subset_size': 300\n",
        "        }\n",
        "    }\n",
        "]\n",
        "\n",
        "# Steps 1-7: Model optimization\n",
        "# Step 1: Load the model.\n",
        "model = load_model(model_config)\n",
        "\n",
        "# Step 2: Initialize the data loader.\n",
        "data_loader = CifarDataLoader(dataset_config)\n",
        "\n",
        "# Step 3 (Optional. Required for AccuracyAwareQuantization): Initialize the metric.\n",
        "metric = Accuracy(top_k=1)\n",
        "\n",
        "# Step 4: Initialize the engine for metric calculation and statistics collection.\n",
        "engine = IEEngine(engine_config, data_loader, metric)\n",
        "\n",
        "# Step 5: Create a pipeline of compression algorithms.\n",
        "pipeline = create_pipeline(algorithms, engine)\n",
        "\n",
        "# Step 6: Execute the pipeline.\n",
        "compressed_model = pipeline.run(model)\n",
        "\n",
        "# Step 7 (Optional): Compress model weights quantized precision\n",
        "#                    in order to reduce the size of final .bin file.\n",
        "compress_model_weights(compressed_model)\n",
        "\n",
        "# Step 8: Save the compressed model to the desired path.\n",
        "compressed_model_paths = save_model(model=compressed_model, save_path=MODEL_DIR, model_name=\"quantized_mobilenet_v2\"\n",
        ")\n",
        "compressed_model_xml = compressed_model_paths[0][\"model\"]\n",
        "compressed_model_bin = Path(compressed_model_paths[0][\"model\"]).with_suffix(\".bin\")\n",
        "\n",
        "# Step 9: Compare accuracy of the original and quantized models.\n",
        "metric_results = pipeline.evaluate(model)\n",
        "if metric_results:\n",
        "    for name, value in metric_results.items():\n",
        "        print(f\"Accuracy of the original model: {name}: {value}\")\n",
        "\n",
        "metric_results = pipeline.evaluate(compressed_model)\n",
        "if metric_results:\n",
        "    for name, value in metric_results.items():\n",
        "        print(f\"Accuracy of the optimized model: {name}: {value}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vQACMfAUo52V"
      },
      "source": [
        "## Compare Performance of the Original and Quantized Models\n",
        "\n",
        "Finally, we will measure the inference performance of the FP32 and INT8 models. To do this, we use [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) - OpenVINO's inference performance measurement tool.\n",
        "\n",
        "NOTE: For more accurate performance, we recommended running 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.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pC0gnO0c9-tI"
      },
      "outputs": [],
      "source": [
        "# Inference FP16 model (IR)\n",
        "!benchmark_app -m $ir_model_xml -d CPU -api async"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4VR3-joFu9hH"
      },
      "outputs": [],
      "source": [
        "# Inference INT8 model (IR)\n",
        "!benchmark_app -m $compressed_model_xml -d CPU -api async"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compare results on four pictures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "ie = Core()\n",
        "\n",
        "# read and load float model\n",
        "float_model = ie.read_model(\n",
        "    model=ir_model_xml, weights=ir_model_bin\n",
        ")\n",
        "float_compiled_model = ie.compile_model(model=float_model, device_name=\"CPU\")\n",
        "\n",
        "# read and load quantized model\n",
        "quantized_model = ie.read_model(\n",
        "    model=compressed_model_xml, weights=compressed_model_bin\n",
        ")\n",
        "quantized_compiled_model = ie.compile_model(model=quantized_model, device_name=\"CPU\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# define all possible labels from CIFAR10\n",
        "labels_names = [\"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\", \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"]\n",
        "all_pictures = []\n",
        "all_labels = []\n",
        "\n",
        "# get all pictures and their labels \n",
        "for i, batch in enumerate(data_loader):\n",
        "    all_pictures.append(batch[1])\n",
        "    all_labels.append(batch[0][1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def plot_pictures(indexes: list, all_pictures=all_pictures, all_labels=all_labels):\n",
        "    \"\"\"Plot 4 pictures.\n",
        "    :param indexes: a list of indexes of pictures to be displayed.\n",
        "    :param all_batches: batches with pictures.\n",
        "    \"\"\"\n",
        "    images, labels = [], []\n",
        "    num_pics = len(indexes)\n",
        "    assert num_pics == 4, f'No enough indexes for pictures to be displayed, got {num_pics}'\n",
        "    for idx in indexes:\n",
        "        assert idx < 10000, 'Cannot get such index, there are only 10000'\n",
        "        pic = np.rollaxis(all_pictures[idx].squeeze(), 0, 3)\n",
        "        images.append(pic)\n",
        "\n",
        "        labels.append(labels_names[all_labels[idx]])\n",
        "\n",
        "    f, axarr = plt.subplots(1, 4)\n",
        "    axarr[0].imshow(images[0])\n",
        "    axarr[0].set_title(labels[0])\n",
        "\n",
        "    axarr[1].imshow(images[1])\n",
        "    axarr[1].set_title(labels[1])\n",
        "\n",
        "    axarr[2].imshow(images[2])\n",
        "    axarr[2].set_title(labels[2])\n",
        "\n",
        "    axarr[3].imshow(images[3])\n",
        "    axarr[3].set_title(labels[3])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def infer_on_pictures(model, indexes: list, all_pictures=all_pictures):\n",
        "    \"\"\" Inference model on a few pictures.\n",
        "    :param net: model on which do inference\n",
        "    :param indexes: list of indexes \n",
        "    \"\"\"\n",
        "    output_key = model.output(0)\n",
        "    predicted_labels = []\n",
        "    for idx in indexes:\n",
        "        assert idx < 10000, 'Cannot get such index, there are only 10000'\n",
        "        result = model([all_pictures[idx][None,]])[output_key]\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 = [7, 12, 15, 20]  # to plot specify 4 indexes\n",
        "\n",
        "plot_pictures(indexes_to_infer)\n",
        "\n",
        "results_float = infer_on_pictures(float_compiled_model, indexes_to_infer)\n",
        "results_quanized = infer_on_pictures(quantized_compiled_model, indexes_to_infer)\n",
        "\n",
        "print(f\"Labels for picture from float model : {results_float}.\")\n",
        "print(f\"Labels for picture from quantized model : {results_quanized}.\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tutorial_tiny.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}