notebooks/102-pytorch-onnx-to-openvino/102-pytorch-onnx-to-openvino.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "af86d01b",
   "metadata": {
    "id": "JwEAhQVzkAwA"
   },
   "source": [
    "# Convert a PyTorch Model to ONNX and OpenVINO™ IR\n",
    "\n",
    "This tutorial demonstrates step-by-step instructions on how to do inference on a PyTorch semantic segmentation model, using OpenVINO Runtime.\n",
    "\n",
    "First, the PyTorch model is converted to [ONNX](https://onnx.ai/) and OpenVINO Intermediate Representation (OpenVINO IR) formats. Then the ONNX and OpenVINO IR models are loaded in OpenVINO Runtime to show model predictions. The model is pre-trained on the [CityScapes](https://www.cityscapes-dataset.com) dataset. The source of the model is [FastSeg](https://github.com/ekzhang/fastseg)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cdf8008",
   "metadata": {
    "id": "QB4Yo-rGGLmV"
   },
   "source": [
    "## Preparation\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2573d828",
   "metadata": {
    "id": "2ynWRum4iiTz"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "from IPython.display import Markdown, display\n",
    "from fastseg import MobileV3Large\n",
    "from openvino.runtime import Core\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from notebook_utils import CityScapesSegmentation, segmentation_map_to_image, viz_result_image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b80fe4",
   "metadata": {},
   "source": [
    "### Settings\n",
    "\n",
    "Set the name for the model, and the width and height of the image that will be used for the network. The CityScapes is pre-trained on images of 2048x1024 resolution. Using smaller dimensions will impact model accuracy, but will improve inference speed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a9303a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_WIDTH = 1024  # Suggested values: 2048, 1024 or 512. The minimum width is 512.\n",
    "# Set IMAGE_HEIGHT manually for custom input sizes. Minimum height is 512.\n",
    "IMAGE_HEIGHT = 1024 if IMAGE_WIDTH == 2048 else 512\n",
    "DIRECTORY_NAME = \"model\"\n",
    "BASE_MODEL_NAME = DIRECTORY_NAME + f\"/fastseg{IMAGE_WIDTH}\"\n",
    "\n",
    "# Paths where PyTorch, ONNX and OpenVINO IR models will be stored.\n",
    "model_path = Path(BASE_MODEL_NAME).with_suffix(\".pth\")\n",
    "onnx_path = model_path.with_suffix(\".onnx\")\n",
    "ir_path = model_path.with_suffix(\".xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b33033e",
   "metadata": {
    "id": "u5xKw0hR0jq6"
   },
   "source": [
    "### Download the Fastseg Model\n",
    "\n",
    "Download, load and save the model with pre-trained weights. This may take some time if you have not downloaded the model before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9600481",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xGKkMRfvi0op",
    "outputId": "4eb1f9af-a4c5-424c-f808-dd9cc2600975"
   },
   "outputs": [],
   "source": [
    "print(\"Downloading the Fastseg model (if it has not been downloaded before)....\")\n",
    "model = MobileV3Large.from_pretrained().cpu().eval()\n",
    "print(\"Loaded PyTorch Fastseg model\")\n",
    "\n",
    "# Save the model.\n",
    "model_path.parent.mkdir(exist_ok=True)\n",
    "torch.save(model.state_dict(), str(model_path))\n",
    "print(f\"Model saved at {model_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bad92bb9",
   "metadata": {
    "id": "Rhc_7EObUypw"
   },
   "source": [
    "## ONNX Model Conversion\n",
    "\n",
    "### Convert PyTorch model to ONNX\n",
    "\n",
    "The output for this cell will show some warnings. These are most likely harmless. When the conversion is successful, the last line of the output will read: `ONNX model exported to fastseg1024.onnx.` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "659aeac7",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ipQWpbgQUxoo",
    "outputId": "bbc1734a-c2a2-4261-ed45-264b9e3edd00"
   },
   "outputs": [],
   "source": [
    "if not onnx_path.exists():\n",
    "    dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH)\n",
    "\n",
    "    # For the Fastseg model, setting do_constant_folding to False is required\n",
    "    # for PyTorch>1.5.1\n",
    "    torch.onnx.export(\n",
    "        model,\n",
    "        dummy_input,\n",
    "        onnx_path,\n",
    "        opset_version=11,\n",
    "        do_constant_folding=False,\n",
    "    )\n",
    "    print(f\"ONNX model exported to {onnx_path}.\")\n",
    "else:\n",
    "    print(f\"ONNX model {onnx_path} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b490241b",
   "metadata": {
    "id": "6JSoEIk60uxV"
   },
   "source": [
    "### Convert ONNX Model to OpenVINO IR Format\n",
    "\n",
    "Use Model Optimizer to convert the ONNX model to OpenVINO IR with `FP16` precision. The models are saved to the current directory. Add mean values to the model and scale the output with the standard deviation with  `--scale_values`. With these options, it is not necessary to normalize input data before propagating it through the network.\n",
    "\n",
    "For more information about Model Optimizer, see the [Model Optimizer Developer Guide](https://docs.openvino.ai/latest/openvino_docs_MO_DG_Deep_Learning_Model_Optimizer_DevGuide.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f99ea10",
   "metadata": {},
   "source": [
    "Executing this command may take a while. There may be some errors or warnings in the output. When Model Optimization is successful, the last lines of the output will include: `[ SUCCESS ] Generated IR version 11 model.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c1e8029",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Construct the command for Model Optimizer.\n",
    "mo_command = f\"\"\"mo\n",
    "                 --input_model \"{onnx_path}\"\n",
    "                 --input_shape \"[1,3, {IMAGE_HEIGHT}, {IMAGE_WIDTH}]\"\n",
    "                 --mean_values=\"[123.675, 116.28 , 103.53]\"\n",
    "                 --scale_values=\"[58.395, 57.12 , 57.375]\"\n",
    "                 --data_type FP16\n",
    "                 --output_dir \"{model_path.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,
   "id": "e48d30a7",
   "metadata": {
    "id": "6YUwrq7QWSzw",
    "tags": []
   },
   "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",
   "id": "608279bb",
   "metadata": {
    "id": "FAGmlKQ83ecE"
   },
   "source": [
    "## Show Results\n",
    "\n",
    "Confirm that the segmentation results look as expected by comparing model predictions on the ONNX, OpenVINO IR and PyTorch models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d2a742",
   "metadata": {
    "id": "JyD5EKka34Wd"
   },
   "source": [
    "### Load and Preprocess an Input Image\n",
    "\n",
    "For the OpenVINO IR model, normalization is moved to the model. For the ONNX and PyTorch models, images need to be normalized before propagating through the network. A sample image from the [Mapillary Vistas](https://www.mapillary.com/dataset/vistas) dataset is provided for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "738cb5e7",
   "metadata": {
    "id": "QTOoQnSetzQM",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def normalize(image: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Normalize the image to the given mean and standard deviation\n",
    "    for CityScapes models.\n",
    "    \"\"\"\n",
    "    image = image.astype(np.float32)\n",
    "    mean = (0.485, 0.456, 0.406)\n",
    "    std = (0.229, 0.224, 0.225)\n",
    "    image /= 255.0\n",
    "    image -= mean\n",
    "    image /= std\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b92da2c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DGFW5VXL3x9G",
    "outputId": "300eacff-c6de-4eb5-e99a-8def5260da1a",
    "tags": []
   },
   "outputs": [],
   "source": [
    "image_filename = \"data/street.jpg\"\n",
    "image = cv2.cvtColor(cv2.imread(image_filename), cv2.COLOR_BGR2RGB)\n",
    "\n",
    "resized_image = cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT))\n",
    "normalized_image = normalize(resized_image)\n",
    "\n",
    "# Convert the resized images to network input shape.\n",
    "input_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)\n",
    "normalized_input_image = np.expand_dims(np.transpose(normalized_image, (2, 0, 1)), 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57bab671",
   "metadata": {
    "id": "FnEiEbNq4Csh"
   },
   "source": [
    "### Load the OpenVINO IR Network and Run Inference on the ONNX model\n",
    "\n",
    "OpenVINO Runtime can load ONNX models directly. First, load the ONNX model, do inference and show the results. Then, load the model that was converted to OpenVINO Intermediate Representation (OpenVINO IR) with Model Optimizer and do inference on that model, and show the results on an image from [Mapillary Vistas](https://www.mapillary.com/dataset/vistas).\n",
    "\n",
    "#### 1. ONNX Model in OpenVINO Runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "194b8bc2",
   "metadata": {
    "id": "otfT6EDk03KV",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the network to OpenVINO Runtime.\n",
    "ie = Core()\n",
    "model_onnx = ie.read_model(model=onnx_path)\n",
    "compiled_model_onnx = ie.compile_model(model=model_onnx, device_name=\"CPU\")\n",
    "\n",
    "output_layer_onnx = compiled_model_onnx.output(0)\n",
    "\n",
    "# Run inference on the input image.\n",
    "res_onnx = compiled_model_onnx([normalized_input_image])[output_layer_onnx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "064a14ba",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 348
    },
    "id": "q8NRH8hLuWeV",
    "outputId": "8b17f90e-b3bc-456f-b6f3-e9cd4f743b2f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Convert the network result to a segmentation map and display the result.\n",
    "result_mask_onnx = np.squeeze(np.argmax(res_onnx, axis=1)).astype(np.uint8)\n",
    "viz_result_image(\n",
    "    image,\n",
    "    segmentation_map_to_image(result_mask_onnx, CityScapesSegmentation.get_colormap()),\n",
    "    resize=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70c0e7d8",
   "metadata": {
    "id": "FnEiEbNq4Csh"
   },
   "source": [
    "#### 2. OpenVINO IR Model in OpenVINO Runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b628c3a4",
   "metadata": {
    "id": "otfT6EDk03KV",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the network in OpenVINO Runtime.\n",
    "ie = Core()\n",
    "model_ir = ie.read_model(model=ir_path)\n",
    "compiled_model_ir = ie.compile_model(model=model_ir, device_name=\"CPU\")\n",
    "\n",
    "# Get input and output layers.\n",
    "output_layer_ir = compiled_model_ir.output(0)\n",
    "\n",
    "# Run inference on the input image.\n",
    "res_ir = compiled_model_ir([input_image])[output_layer_ir]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e625d64a",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 348
    },
    "id": "q8NRH8hLuWeV",
    "outputId": "8b17f90e-b3bc-456f-b6f3-e9cd4f743b2f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_mask_ir = np.squeeze(np.argmax(res_ir, axis=1)).astype(np.uint8)\n",
    "viz_result_image(\n",
    "    image,\n",
    "    segmentation_map_to_image(result=result_mask_ir, colormap=CityScapesSegmentation.get_colormap()),\n",
    "    resize=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b8de873",
   "metadata": {
    "id": "w3UUduQEGsQm"
   },
   "source": [
    "## PyTorch Comparison\n",
    "\n",
    "Do inference on the PyTorch model to verify that the output visually looks the same as the output on the ONNX/OpenVINO IR models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b4f32b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 348
    },
    "id": "1l1JtgeV4Wuw",
    "outputId": "f21c8904-83da-438c-df39-4620bb679554",
    "tags": []
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    result_torch = model(torch.as_tensor(normalized_input_image).float())\n",
    "\n",
    "result_mask_torch = torch.argmax(result_torch, dim=1).squeeze(0).numpy().astype(np.uint8)\n",
    "viz_result_image(\n",
    "    image,\n",
    "    segmentation_map_to_image(result=result_mask_torch, colormap=CityScapesSegmentation.get_colormap()),\n",
    "    resize=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9feb7b4",
   "metadata": {},
   "source": [
    "## Performance Comparison\n",
    "\n",
    "Measure the time it takes to do inference on twenty images. This gives an indication of performance. For more accurate benchmarking, use the [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html). Keep in mind that many optimizations are possible to improve the performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "835b75e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_images = 20\n",
    "\n",
    "with torch.no_grad():\n",
    "    start = time.perf_counter()\n",
    "    for _ in range(num_images):\n",
    "        model(torch.as_tensor(input_image).float())\n",
    "    end = time.perf_counter()\n",
    "    time_torch = end - start\n",
    "print(\n",
    "    f\"PyTorch model on CPU: {time_torch/num_images:.3f} seconds per image, \"\n",
    "    f\"FPS: {num_images/time_torch:.2f}\"\n",
    ")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for _ in range(num_images):\n",
    "    compiled_model_onnx([normalized_input_image])\n",
    "end = time.perf_counter()\n",
    "time_onnx = end - start\n",
    "print(\n",
    "    f\"ONNX model in OpenVINO Runtime/CPU: {time_onnx/num_images:.3f} \"\n",
    "    f\"seconds per image, FPS: {num_images/time_onnx:.2f}\"\n",
    ")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for _ in range(num_images):\n",
    "    compiled_model_ir([input_image])\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(\n",
    "    f\"OpenVINO IR model in OpenVINO Runtime/CPU: {time_ir/num_images:.3f} \"\n",
    "    f\"seconds per image, FPS: {num_images/time_ir:.2f}\"\n",
    ")\n",
    "\n",
    "if \"GPU\" in ie.available_devices:\n",
    "    compiled_model_onnx_gpu = ie.compile_model(model=model_onnx, device_name=\"GPU\")\n",
    "    start = time.perf_counter()\n",
    "    for _ in range(num_images):\n",
    "        compiled_model_onnx_gpu([input_image])\n",
    "    end = time.perf_counter()\n",
    "    time_onnx_gpu = end - start\n",
    "    print(\n",
    "        f\"ONNX model in OpenVINO/GPU: {time_onnx_gpu/num_images:.3f} \"\n",
    "        f\"seconds per image, FPS: {num_images/time_onnx_gpu:.2f}\"\n",
    "    )\n",
    "\n",
    "    compiled_model_ir_gpu = ie.compile_model(model=model_ir, device_name=\"GPU\")\n",
    "    start = time.perf_counter()\n",
    "    for _ in range(num_images):\n",
    "        compiled_model_ir_gpu([input_image])\n",
    "    end = time.perf_counter()\n",
    "    time_ir_gpu = end - start\n",
    "    print(\n",
    "        f\"IR model in OpenVINO/GPU: {time_ir_gpu/num_images:.3f} \"\n",
    "        f\"seconds per image, FPS: {num_images/time_ir_gpu:.2f}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6ca50ca",
   "metadata": {},
   "source": [
    "**Show Device Information**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab884bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "devices = ie.available_devices\n",
    "for device in devices:\n",
    "    device_name = ie.get_property(device, \"FULL_DEVICE_NAME\")\n",
    "    print(f\"{device}: {device_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8ed2610",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* [Fastseg](https://github.com/ekzhang/fastseg)\n",
    "* [PIP install openvino-dev](https://github.com/openvinotoolkit/openvino/blob/releases/2021/3/docs/install_guides/pypi-openvino-dev.md)\n",
    "* [OpenVINO ONNX support](https://docs.openvino.ai/2021.4/openvino_docs_IE_DG_ONNX_Support.html)\n",
    "* [Model Optimizer Documentation](https://docs.openvino.ai/latest/openvino_docs_MO_DG_prepare_model_convert_model_Converting_Model_General.html)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "ae617ccb002f72b3ab6d0069d721eac67ac2a969e83c083c4321cfcab0437cd1"
  },
  "kernelspec": {
   "display_name": "openvino_env",
   "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}