demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "be2cc57c",
   "metadata": {},
   "source": [
    "## Evaluate baseline model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba199ca",
   "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": "5166c0c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from val import run as run_val\n",
    "\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'yolov5s.pt',\n",
    "    'half': True,\n",
    "    'batch_size': 3, # 32\n",
    "}\n",
    "\n",
    "run_val(**opt);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "085c9710",
   "metadata": {},
   "source": [
    "## Optimize model with ENOT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7697bfbb-37b8-4c24-abcb-8c999a073a9f",
   "metadata": {},
   "source": [
    "Metrics after optimization may be very low, but it will recover after tuning phase. Depending on `target-latency-fraction`, drop of the metrics can be different. The complexity of your detection task also affects possible drop of the metrics after pruning.\n",
    "\n",
    "It is always the trade-off between how fast you want your model to be and the final metrics after fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58eaeb32",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from prune import run as run_prune\n",
    "\n",
    "opt = {\n",
    "    'device': 0,\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'yolov5s.pt',\n",
    "    'half': True,\n",
    "    'batch_size': 3, # Better use batch_size > 8 because of batch norm\n",
    "    'imgsz': 640,\n",
    "    'hyp': 'data/hyps/hyp.coco_pruning.yaml',\n",
    "    'name': 'prune_yolov5s_coco',\n",
    "    'save_before_prune': True,\n",
    "    'n_search_steps': 3, # This value is just for demo, in production we recommend to use more than 200 steps.\n",
    "    'target-latency-fraction': 0.5, # It means that optimized model will be 2 times faster than baseline.\n",
    "}\n",
    "\n",
    "run_prune(**opt);\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed2ac368",
   "metadata": {},
   "source": [
    "## Make onnx for original and optimized models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e366058-95f9-4b0f-b630-df0e782e8a9d",
   "metadata": {},
   "source": [
    "You can open ONNXes using [Netron](https://netron.app/) and visualize model architectures. You can see that pruned model has different number of channels compared to the original one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e762bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from export import run as run_export\n",
    "\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'runs/prune/prune_yolov5s_coco/weights/original_model.pt',\n",
    "    'batch_size': 1,\n",
    "    'imgsz': [640],\n",
    "    'include': ['onnx'],\n",
    "}\n",
    "\n",
    "run_export(**opt)\n",
    "\n",
    "opt['weights'] = 'runs/prune/prune_yolov5s_coco/weights/pruned_model.pt'\n",
    "\n",
    "run_export(**opt)\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9eee885",
   "metadata": {},
   "source": [
    "## Run optimized model tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d119bbed-b4a3-4015-9a42-ef9f5ce42c19",
   "metadata": {},
   "source": [
    "This phase is required because just after pruning with low `target-latency-fraction` metrics will be close to zero. But just after the first epoch of tuning metrics will start to recover."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "194b6b27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from train import run as run_tune\n",
    "\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'runs/prune/prune_yolov5s_coco/weights/pruned_model.pt',\n",
    "    'batch_size': 3, # Better use batch_size > 8 because of batch norm\n",
    "    'imgsz': 640,\n",
    "    'from_pruned': True,\n",
    "    'epochs': 1, # Just for demo, in the real task it needs to be a half of baseline training epochs or lower.\n",
    "    'device': 0,\n",
    "    'name': 'tune_pruned_model',\n",
    "}\n",
    "\n",
    "run_tune(**opt)\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80c7ea6d",
   "metadata": {},
   "source": [
    "## Evaluate tuned optimized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "774184d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from val import run as run_val\n",
    "\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'runs/train/tune_pruned_model/weights/best.pt',\n",
    "    'half': True,\n",
    "    'batch_size': 3, # Better use batch_size > 8 because of batch norm\n",
    "    'imgsz': 640,\n",
    "}\n",
    "\n",
    "run_val(**opt);\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9dd3d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from detect import run as run_detect\n",
    "\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'source': '../datasets/coco128/images/train2017/',\n",
    "    'weights': 'runs/train/tune_pruned_model/weights/best.pt',\n",
    "    'half': True,\n",
    "    'imgsz': (640, 640),\n",
    "    'name': 'optimized_model'\n",
    "}\n",
    "run_detect(**opt)\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "opt['name'] = 'original_model'\n",
    "opt['weights'] = 'yolov5s.pt'\n",
    "run_detect(**opt)\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bcf0558",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show detection results for original and optimized models\n",
    "\n",
    "%matplotlib inline\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "original_predict = cv2.imread('runs/detect/original_model/000000000009.jpg')\n",
    "optimized_predict = cv2.imread('runs/detect/optimized_model/000000000009.jpg')\n",
    "\n",
    "figsize = 10\n",
    "plt.figure(figsize=(figsize, figsize))\n",
    "plt.imshow(cv2.hconcat([original_predict, optimized_predict])[:,:,::-1])\n",
    "plt.grid(visible=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb253f6a",
   "metadata": {},
   "source": [
    "# OpenVino quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f89ee8ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from export import run as run_export\n",
    "\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'runs/train/tune_pruned_model/weights/best.pt',\n",
    "    'batch_size': 1,\n",
    "    'imgsz': [640],\n",
    "    'include': ['onnx'],\n",
    "}\n",
    "\n",
    "run_export(**opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b5dcbdf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=0\n",
    "from quant import run as run_quant\n",
    "opt = {\n",
    "    'data': 'data/coco128.yaml',\n",
    "    'weights': 'runs/train/tune_pruned_model/weights/best.onnx',\n",
    "    'batch_size': 1,\n",
    "    'imgsz': 640,\n",
    "    'device': 'cuda',\n",
    "    'backend': 'openvino',\n",
    "    'n_epochs': 2,\n",
    "}\n",
    "\n",
    "run_quant(**opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0524eaf3",
   "metadata": {},
   "source": [
    "# Run quantized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643f857b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from enot_lite.backend import BackendFactory\n",
    "from enot_lite.type import BackendType\n",
    "\n",
    "inputs = np.ones((1,3,640,640), dtype=np.float32)\n",
    "backend = BackendFactory().create(\n",
    "    'runs/train/tune_pruned_model/weights/best_quant.onnx',\n",
    "    BackendType.ORT_OPENVINO,\n",
    "    input_example=inputs,\n",
    ")\n",
    "\n",
    "prediction = backend(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bbb2087",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1bbe923",
   "metadata": {},
   "source": [
    "# Check acceleration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ccd79ff",
   "metadata": {},
   "source": [
    "### Baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c4ea1e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from enot_lite.benchmark import Benchmark\n",
    "from enot_lite.type import BackendType\n",
    "\n",
    "benchmark = Benchmark(\n",
    "    batch_size=1,\n",
    "    onnx_model='runs/prune/prune_yolov5s_coco/weights/original_model.onnx',\n",
    "    onnx_input=(np.ones((1, 3, 640, 640), dtype=np.float32),),\n",
    "    backends=[BackendType.ORT_OPENVINO],\n",
    "    number=10,\n",
    "    warmup=10,\n",
    "    repeat=10\n",
    ")\n",
    "\n",
    "benchmark.run()\n",
    "benchmark.print_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ccaef57",
   "metadata": {},
   "source": [
    "### Pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d5c9fdf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from enot_lite.benchmark import Benchmark\n",
    "from enot_lite.type import BackendType\n",
    "\n",
    "benchmark = Benchmark(\n",
    "    batch_size=1,\n",
    "    onnx_model='runs/train/tune_pruned_model/weights/best.onnx',\n",
    "    onnx_input=(np.ones((1, 3, 640, 640), dtype=np.float32),),\n",
    "    backends=[BackendType.ORT_OPENVINO],\n",
    "    number=10,\n",
    "    warmup=10,\n",
    "    repeat=10\n",
    ")\n",
    "\n",
    "benchmark.run()\n",
    "benchmark.print_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9f31a5b",
   "metadata": {},
   "source": [
    "### Quantized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07c949a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from enot_lite.benchmark import Benchmark\n",
    "from enot_lite.type import BackendType\n",
    "\n",
    "benchmark = Benchmark(\n",
    "    batch_size=1,\n",
    "    onnx_model='runs/train/tune_pruned_model/weights/best_quant.onnx',\n",
    "    onnx_input=(np.ones((1, 3, 640, 640), dtype=np.float32),),\n",
    "    backends=[BackendType.ORT_OPENVINO],\n",
    "    number=10,\n",
    "    warmup=10,\n",
    "    repeat=10,\n",
    ")\n",
    "\n",
    "benchmark.run()\n",
    "benchmark.print_results()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8341e5cf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}