3. Tutorial - Ultralytics YOLOv8 quantization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "826e0457",
   "metadata": {},
   "source": [
    "## Automatic quantization for YOLOv8\n",
    "\n",
    "This notebook demonstrates simple procedure for Ultralytics YOLOv8 quantization for OpenVINO.\n",
    "\n",
    "Our quantization process consists of quantized model calibration, quantization thresholds adjustment and weight fine-tuning using distillation. Finally, we demonstrate inference of our quantized model using YOLOv8 and OpenVINO frameworks.\n",
    "\n",
    "### Main chapters of this notebook:\n",
    "1. Setup the environment\n",
    "2. Prepare dataset and create dataloaders\n",
    "3. Export YOLOv8 to ONNX\n",
    "4. Quantize YOLOv8\n",
    "5. Measure inference time using OpenVINO framework\n",
    "6. Measure mAP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a319920",
   "metadata": {},
   "source": [
    "## Setup the environment\n",
    "\n",
    "First, let's set up the environment and make some common imports."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c6e4737",
   "metadata": {},
   "source": [
    "1. Install `enot-autodl` package and create Jupyter kernel with it\n",
    "2. Install `ultralytics` package with YOLOv8\n",
    "3. Install `openvino` and `openvino-dev`\n",
    "\n",
    "To install `enot-autodl` package follow the [installation guide](https://enot-autodl.rtd.enot.ai/en/latest/installation_guide.html).  \n",
    "For p. 2-3 see commands below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "187eb985",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install ultralytics==8.0.199\n",
    "!pip install openvino==2023.2.0 openvino-dev==2023.2.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1440218",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You may need to uncomment and change this variable to match free GPU index\n",
    "# %env CUDA_VISIBLE_DEVICES=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d93b61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "import itertools\n",
    "import numpy as np\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import torch\n",
    "from torch.optim import RAdam\n",
    "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
    "\n",
    "# quantization procedure\n",
    "from enot.quantization import distill\n",
    "from enot.quantization import OpenVINOFakeQuantizedModel\n",
    "from enot.quantization import calibrate\n",
    "from enot.quantization import RMSELoss\n",
    "\n",
    "# converters from onnx to pytorch\n",
    "from onnx2torch import convert\n",
    "\n",
    "# dataset creation functions\n",
    "from ultralytics.utils import DEFAULT_CFG\n",
    "from ultralytics.cfg import get_cfg\n",
    "from ultralytics.data.utils import check_det_dataset\n",
    "from ultralytics.data import build_dataloader, build_yolo_dataset\n",
    "\n",
    "# function for loading yolo checkpoint\n",
    "from ultralytics import YOLO\n",
    "\n",
    "# openvino functions\n",
    "from tutorial_utils.openvino import benchmark\n",
    "from tutorial_utils.openvino import convert_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4eb6f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "QUANT_ONNX_PATH = './yolov8s_openvino_int8.onnx'\n",
    "\n",
    "OV_FP32_NAME = \"yolov8s_fp32\"\n",
    "OV_INT8_NAME = \"yolov8s_int8\"\n",
    "\n",
    "OV_FP32_FULL_NAME = f\"{OV_FP32_NAME}_openvino_model/{OV_FP32_NAME}.xml\"\n",
    "OV_INT8_FULL_NAME = f\"{OV_INT8_NAME}_openvino_model/{OV_INT8_NAME}.xml\"\n",
    "\n",
    "BATCH_SIZE = 8\n",
    "IMG_SIZE = 640\n",
    "IMG_SHAPE = (BATCH_SIZE, 3, IMG_SIZE, IMG_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f2ec45",
   "metadata": {},
   "source": [
    "## Prepare dataset and create dataloaders\n",
    "\n",
    "We will use MS COCO128 dataset in this example.\n",
    "\n",
    "\n",
    "`build_dataloader`, `check_det_dataset` and `build_yolo_dataset` functions prepare datasets for you in this example; specifically, it:\n",
    "1. Downloads and unpacks dataset to `datasets/coco128` or to existing YOLOv8 data path if `ultralytics` package is installed\n",
    "2. Creates and returns train and validation dataloaders\n",
    "\n",
    "**IMPORTANT NOTE**: since this is example notebook we will train and validate model in **THE SAME DATASET**. For better performance and generalization use separate dataset for train and val procedure. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "194cc18b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = get_cfg(DEFAULT_CFG, None)\n",
    "cfg.data = 'coco128.yaml'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "308fe28f",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = check_det_dataset(cfg.data)\n",
    "if 'yaml_file' in data:\n",
    "    cfg.data = data['yaml_file']\n",
    "\n",
    "trainset, testset = data['train'], data.get('val') or data.get('test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "191b2c71",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = build_yolo_dataset(\n",
    "    cfg=cfg,\n",
    "    img_path=trainset,\n",
    "    batch=BATCH_SIZE,\n",
    "    data=data,\n",
    ")\n",
    "\n",
    "dataloader = build_dataloader(\n",
    "    dataset=dataset,\n",
    "    batch=BATCH_SIZE,\n",
    "    workers=cfg.workers,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e874892",
   "metadata": {},
   "source": [
    "## Baseline YOLOv8 ONNX creation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b6da6cb",
   "metadata": {},
   "source": [
    "Since the default YOLOv8 model contains conditional execution ('if' nodes), we have to save it to ONNX format and convert back to PyTorch to perform quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6d52436",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLO(model='yolov8s')\n",
    "onnx_path = model.export(format='onnx', dynamic=True, imgsz=IMG_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5837bfc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "regular_model = convert(onnx_path).cuda()\n",
    "regular_model.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43d4eea0",
   "metadata": {},
   "source": [
    "## YOLOv8 Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9be645c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's define function for converting dataset samples to model inputs.\n",
    "\n",
    "\n",
    "def sample_to_model_inputs(x):\n",
    "    # x[0] is the first item from dataloader sample. Sample is a tuple where 0'th element is a tensor with images.\n",
    "    x = x['img']\n",
    "\n",
    "    # Model is on CUDA, so input images should also be on CUDA.\n",
    "    x = x.cuda()\n",
    "\n",
    "    # Converting tensor from int8 to float data type.\n",
    "    x = x.float()\n",
    "\n",
    "    # YOLOv8 image normalization (0-255 to 0-1 normalization)\n",
    "    x /= 255\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f31e8835-aec8-458c-bc26-05239374b38d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See for details: https://enot-autodl.rtd.enot.ai/en/stable/reference_documentation/quantization.html#enot.quantization.OpenVINOFakeQuantizedModel\n",
    "\n",
    "fake_quantized_model = OpenVINOFakeQuantizedModel(regular_model).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3604b6d1-afd5-4b18-af92-7e99c935d58d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calibrate quantization thresholds using 10 batches.\n",
    "\n",
    "with torch.no_grad(), calibrate(fake_quantized_model):\n",
    "    for batch in itertools.islice(dataloader, 10):\n",
    "        batch = sample_to_model_inputs(batch)\n",
    "        fake_quantized_model(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "463e7dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Distill model quantization thresholds and weights using RMSE loss.\n",
    "\n",
    "n_epochs = 5\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(dataloader) * n_epochs)\n",
    "    distillation_criterion = RMSELoss()\n",
    "\n",
    "    for _ in range(n_epochs):\n",
    "        for batch in (tqdm_it := tqdm(dataloader)):\n",
    "            batch = sample_to_model_inputs(batch)\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,
   "id": "008bb58c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fake_quantized_model.cuda()\n",
    "fake_quantized_model.enable_quantization_mode(True)\n",
    "fake_quantized_model.cpu()\n",
    "\n",
    "torch.onnx.export(\n",
    "    model=fake_quantized_model,\n",
    "    args=torch.ones(*IMG_SHAPE),\n",
    "    f=QUANT_ONNX_PATH,\n",
    "    input_names=['images'],\n",
    "    output_names=['output'],\n",
    "    opset_version=13,\n",
    "    dynamic_axes={'images': {0: 'batch_size'}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f48f130",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2382f25f",
   "metadata": {},
   "source": [
    "## Measure models speed using OpenVINO framework"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e3207dd",
   "metadata": {},
   "source": [
    "### OpenVINO FP32 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93fbfced",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Convert yolov8s.pt to OpenVINO model\n",
    "yolov8s = YOLO('yolov8s.pt')\n",
    "yolov8s.model.pt_path = OV_FP32_NAME + \".pt\"\n",
    "yolov8s.export(format='openvino', dynamic=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be4b4ca4",
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark(OV_FP32_FULL_NAME, IMG_SHAPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a04800af",
   "metadata": {},
   "source": [
    "### ENOT OpenVINO INT8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d6462c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert yolov8s_enot_int8.onnx to OpenVINO model\n",
    "convert_model(QUANT_ONNX_PATH, OV_INT8_NAME)\n",
    "\n",
    "# copy metadata for YOLO to understand classes and shapes\n",
    "shutil.copy(OV_FP32_NAME + \"_openvino_model/metadata.yaml\", OV_INT8_NAME + \"_openvino_model\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b3583d",
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark(OV_INT8_FULL_NAME, IMG_SHAPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e946a32e",
   "metadata": {},
   "source": [
    "## Measure mAP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45f281b2",
   "metadata": {},
   "source": [
    "### OpenVINO FP32 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf73ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure you have converted OpenVINO model\n",
    "YOLO(OV_FP32_NAME + \"_openvino_model\", task='detect').val(data=cfg.data, imgsz=IMG_SIZE);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4adde54e",
   "metadata": {},
   "source": [
    "### OpenVINO INT8 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e0f5a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure you have converted OpenVINO model\n",
    "YOLO(OV_INT8_NAME + \"_openvino_model\", task='detect').val(data=cfg.data, imgsz=IMG_SIZE);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf48c64f",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 3000
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}