tutorial/hello_world/hello_world.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hello World Example\n",
    "\n",
    "This is a simple Jupyter Notebook that walks through the 4 steps of compiling and running a PyTorch model on the embedded Neural Processing Unit (NPU) in your AMD Ryzen AI enabled PC. The steps are as follows:\n",
    "\n",
    "1. Get model - download or create a PyTorch model that we will run on the NPU\n",
    "2. Export to ONNX - convert the PyTorch model to ONNX format.\n",
    "3. Quantize - optimize the model for faster inference on the NPU by reducing its precision to INT8.\n",
    "4. Run Model on CPU and NPU - compare performance between running the model on the CPU and on the NPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from -r requirements.txt (line 1)) (2.4.0)\n",
      "Requirement already satisfied: ipykernel in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from -r requirements.txt (line 2)) (6.29.5)\n",
      "Requirement already satisfied: filelock in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from torch->-r requirements.txt (line 1)) (3.15.4)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from torch->-r requirements.txt (line 1)) (4.12.2)\n",
      "Requirement already satisfied: sympy in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from torch->-r requirements.txt (line 1)) (1.13.2)\n",
      "Requirement already satisfied: networkx in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from torch->-r requirements.txt (line 1)) (3.3)\n",
      "Requirement already satisfied: jinja2 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from torch->-r requirements.txt (line 1)) (3.1.4)\n",
      "Requirement already satisfied: fsspec in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from torch->-r requirements.txt (line 1)) (2024.6.1)\n",
      "Requirement already satisfied: comm>=0.1.1 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (0.2.2)\n",
      "Requirement already satisfied: debugpy>=1.6.5 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (1.8.5)\n",
      "Requirement already satisfied: ipython>=7.23.1 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (8.26.0)\n",
      "Requirement already satisfied: jupyter-client>=6.1.12 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (8.6.2)\n",
      "Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (5.7.2)\n",
      "Requirement already satisfied: matplotlib-inline>=0.1 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (0.1.7)\n",
      "Requirement already satisfied: nest-asyncio in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (1.6.0)\n",
      "Requirement already satisfied: packaging in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (24.1)\n",
      "Requirement already satisfied: psutil in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (6.0.0)\n",
      "Requirement already satisfied: pyzmq>=24 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (26.1.0)\n",
      "Requirement already satisfied: tornado>=6.1 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (6.4.1)\n",
      "Requirement already satisfied: traitlets>=5.4.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipykernel->-r requirements.txt (line 2)) (5.14.3)\n",
      "Requirement already satisfied: decorator in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (5.1.1)\n",
      "Requirement already satisfied: jedi>=0.16 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (0.19.1)\n",
      "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (3.0.47)\n",
      "Requirement already satisfied: pygments>=2.4.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (2.18.0)\n",
      "Requirement already satisfied: stack-data in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (0.6.2)\n",
      "Requirement already satisfied: exceptiongroup in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (1.2.2)\n",
      "Requirement already satisfied: colorama in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (0.4.6)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from jupyter-client>=6.1.12->ipykernel->-r requirements.txt (line 2)) (2.9.0)\n",
      "Requirement already satisfied: platformdirs>=2.5 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from jupyter-core!=5.0.*,>=4.12->ipykernel->-r requirements.txt (line 2)) (4.2.2)\n",
      "Requirement already satisfied: pywin32>=300 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from jupyter-core!=5.0.*,>=4.12->ipykernel->-r requirements.txt (line 2)) (306)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from jinja2->torch->-r requirements.txt (line 1)) (2.1.5)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from sympy->torch->-r requirements.txt (line 1)) (1.3.0)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.3 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from jedi>=0.16->ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (0.8.4)\n",
      "Requirement already satisfied: wcwidth in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (0.2.13)\n",
      "Requirement already satisfied: six>=1.5 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from python-dateutil>=2.8.2->jupyter-client>=6.1.12->ipykernel->-r requirements.txt (line 2)) (1.16.0)\n",
      "Requirement already satisfied: executing>=1.2.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from stack-data->ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (2.0.1)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from stack-data->ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (2.4.1)\n",
      "Requirement already satisfied: pure-eval in c:\\users\\vgods\\miniconda3\\envs\\ryzen-ai-1.2.0\\lib\\site-packages (from stack-data->ipython>=7.23.1->ipykernel->-r requirements.txt (line 2)) (0.2.3)\n"
     ]
    }
   ],
   "source": [
    "# Before starting, be sure you've installed the requirements listed in the requirements.txt file:\n",
    "!python -m pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0. Imports & Environment Variables\n",
    "\n",
    "We'll use the following imports in our example. `torch` and `torch_nn` are used for building and running ML models. We'll use them to define a small neural network and to generate the model weights. `os` is used for interacting with the operating system and is used to manage our environment variables, file paths, and directories. `subprocess` allows us to retrieve the hardware information. `onnx` and `onnxruntime` are used to work with our model in the ONNX format and for running our inference. `vai_q_onnx` is part of the Vitis AI Quantizer for ONNX models. We use it to perform quantization, converting the model into an INT8 format that is optimized for the NPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import os\n",
    "import subprocess\n",
    "import onnxruntime\n",
    "import numpy as np\n",
    "import onnx\n",
    "import shutil\n",
    "from timeit import default_timer as timer\n",
    "import vai_q_onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As well, we want to set the environment variables based on the NPU device we have in our PC. For more information about NPU configurations, see: For more information about NPU configurations, refer to the official [AMD Ryzen AI Documentation](https://ryzenai.docs.amd.com/en/latest/runtime_setup.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "APU Type: PHX/HPT\n"
     ]
    }
   ],
   "source": [
    "# This function detects the APU (NPU) type in your system to configure environment variables for hardware-specific optimization.\n",
    "def get_apu_info():\n",
    "    # Run pnputil as a subprocess to enumerate PCI devices\n",
    "    command = r'pnputil /enum-devices /bus PCI /deviceids '\n",
    "    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
    "    stdout, stderr = process.communicate()\n",
    "    # Check for supported Hardware IDs\n",
    "    apu_type = ''\n",
    "    if 'PCI\\\\VEN_1022&DEV_1502&REV_00' in stdout.decode(): apu_type = 'PHX/HPT'\n",
    "    if 'PCI\\\\VEN_1022&DEV_17F0&REV_00' in stdout.decode(): apu_type = 'STX'\n",
    "    if 'PCI\\\\VEN_1022&DEV_17F0&REV_10' in stdout.decode(): apu_type = 'STX'\n",
    "    if 'PCI\\\\VEN_1022&DEV_17F0&REV_11' in stdout.decode(): apu_type = 'STX'\n",
    "    return apu_type\n",
    "\n",
    "apu_type = get_apu_info()\n",
    "print(f\"APU Type: {apu_type}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting environment for PHX/HPT\n",
      "XLNX_VART_FIRMWARE= C:\\Program Files\\RyzenAI\\1.2.0\\voe-4.0-win_amd64\\xclbins\\phoenix\\1x4.xclbin\n",
      "NUM_OF_DPU_RUNNERS= 1\n",
      "XLNX_TARGET_NAME= AMD_AIE2_Nx4_Overlay\n"
     ]
    }
   ],
   "source": [
    "# XLNX_VART_FIRMWARE - Specifies the firmware file used by the NPU for runtime execution\n",
    "# NUM_OF_DPU_RUNNERS - Specifies the number of DPU runners (processing cores) available for execution\n",
    "# XLNX_TARGET_NAME - Name of the target hardware configuration\n",
    "\n",
    "def set_environment_variable(apu_type):\n",
    "\n",
    "    install_dir = os.environ['RYZEN_AI_INSTALLATION_PATH']\n",
    "    match apu_type:\n",
    "        case 'PHX/HPT':\n",
    "            print(\"Setting environment for PHX/HPT\")\n",
    "            os.environ['XLNX_VART_FIRMWARE']= os.path.join(install_dir, 'voe-4.0-win_amd64', 'xclbins', 'phoenix', '1x4.xclbin')\n",
    "            os.environ['NUM_OF_DPU_RUNNERS']='1'\n",
    "            os.environ['XLNX_TARGET_NAME']='AMD_AIE2_Nx4_Overlay'\n",
    "        case 'STX':\n",
    "            print(\"Setting environment for STX\")\n",
    "            os.environ['XLNX_VART_FIRMWARE']= os.path.join(install_dir, 'voe-4.0-win_amd64', 'xclbins', 'strix', 'AMD_AIE2P_Nx4_Overlay.xclbin')\n",
    "            os.environ['NUM_OF_DPU_RUNNERS']='1'\n",
    "            os.environ['XLNX_TARGET_NAME']='AMD_AIE2_Nx4_Overlay'\n",
    "        case _:\n",
    "            print(\"Unrecognized APU type. Exiting.\")\n",
    "            exit()\n",
    "    print('XLNX_VART_FIRMWARE=', os.environ['XLNX_VART_FIRMWARE'])\n",
    "    print('NUM_OF_DPU_RUNNERS=', os.environ['NUM_OF_DPU_RUNNERS'])\n",
    "    print('XLNX_TARGET_NAME=', os.environ['XLNX_TARGET_NAME'])\n",
    "\n",
    "set_environment_variable(apu_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Get Model\n",
    "Here, we'll use the PyTorch library to define and instantiate a simple neural network model called `SmallModel` as a starting point. You can swap this model with any custom model, but make sure the input/output shapes remain compatible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SmallModel(\n",
      "  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (conv4): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (relu): ReLU()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "class SmallModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SmallModel, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.relu(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        x = self.relu(x) \n",
    "        \n",
    "        x = self.conv3(x)\n",
    "        x = self.relu(x) \n",
    "        \n",
    "        x = self.conv4(x)\n",
    "        x = self.relu(x) \n",
    "        \n",
    "        x = torch.add(x, 1)\n",
    "        \n",
    "        return x\n",
    "\n",
    "# Instantiate the model\n",
    "pytorch_model = SmallModel()\n",
    "\n",
    "pytorch_model.eval()\n",
    "\n",
    "# Print the model architecture\n",
    "print(pytorch_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Export to ONNX\n",
    "\n",
    "The following code is used for exporting a PyTorch model (pytorch_model) to the ONNX (Open Neural Network Exchange) format. ONNX is an open format that facilitates interoperability between different AI frameworks. Ryzen AI uses ONNX as the input format for quantization using the Vitis AI Quantizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate dummy input data\n",
    "batch_size = 1\n",
    "input_channels = 3\n",
    "input_size = 224\n",
    "dummy_input = torch.rand(batch_size, input_channels, input_size, input_size)\n",
    "\n",
    "# Prep for ONNX export\n",
    "inputs = {\"x\": dummy_input}\n",
    "dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}\n",
    "tmp_model_path = \"models/helloworld.onnx\"\n",
    "\n",
    "# Call export function\n",
    "torch.onnx.export(\n",
    "        pytorch_model,\n",
    "        inputs,\n",
    "        tmp_model_path,\n",
    "        export_params=True,\n",
    "        opset_version=17,  # Recommended opset\n",
    "        input_names=['input'],\n",
    "        output_names=['output'],\n",
    "        dynamic_axes=dynamic_axes,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Quantize Model\n",
    "\n",
    "Using the static quantization method provided by the Vitis AI Quantizer and providing the newly exported ONNX model, we'll quantize the model to INT8. Quantization reduces the precision of model weights and activations from 32-bit floating point (FP32) to 8-bit integers (INT8). This compression allows the model to run faster on hardware accelerators like NPUs, while maintaining nearly the same accuracy. For more information on this quantization method, see [Vitis AI ONNX Quantization](https://ryzenai.docs.amd.com/en/latest/vai_quant/vai_q_onnx.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:vai_q_onnx.quantize:calibration_data_reader is None, using random data for calibration\n",
      "INFO:vai_q_onnx.quant_utils:The input ONNX model models/helloworld.onnx can create InferenceSession successfully\n",
      "INFO:vai_q_onnx.quant_utils:Random input name input shape [1, 3, 224, 224] type <class 'numpy.float32'> \n",
      "INFO:vai_q_onnx.quant_utils:Obtained calibration data with 1 iters\n",
      "INFO:vai_q_onnx.quantize:Removed initializers from input\n",
      "INFO:vai_q_onnx.quantize:Simplified model sucessfully\n",
      "INFO:vai_q_onnx.quantize:Loading model...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[VAI_Q_ONNX_INFO]: Time information:\n",
      "2024-08-23 10:12:35.362481\n",
      "[VAI_Q_ONNX_INFO]: OS and CPU information:\n",
      "                                        system --- Windows\n",
      "                                          node --- vgodsoe-ryzen\n",
      "                                       release --- 10\n",
      "                                       version --- 10.0.26100\n",
      "                                       machine --- AMD64\n",
      "                                     processor --- AMD64 Family 25 Model 116 Stepping 1, AuthenticAMD\n",
      "[VAI_Q_ONNX_INFO]: Tools version information:\n",
      "                                        python --- 3.10.14\n",
      "                                          onnx --- 1.16.2\n",
      "                                   onnxruntime --- 1.17.0\n",
      "                                    vai_q_onnx --- 1.17.0+511d6f4\n",
      "[VAI_Q_ONNX_INFO]: Quantized Configuration information:\n",
      "                                   model_input --- models/helloworld.onnx\n",
      "                                  model_output --- models/helloworld_quantized.onnx\n",
      "                       calibration_data_reader --- None\n",
      "                         calibration_data_path --- None\n",
      "                                  quant_format --- QDQ\n",
      "                                   input_nodes --- []\n",
      "                                  output_nodes --- []\n",
      "                          op_types_to_quantize --- []\n",
      "                random_data_reader_input_shape --- []\n",
      "                                   per_channel --- False\n",
      "                                  reduce_range --- False\n",
      "                               activation_type --- QUInt8\n",
      "                                   weight_type --- QInt8\n",
      "                             nodes_to_quantize --- []\n",
      "                              nodes_to_exclude --- []\n",
      "                                optimize_model --- True\n",
      "                      use_external_data_format --- False\n",
      "                              calibrate_method --- PowerOfTwoMethod.MinMSE\n",
      "                           execution_providers --- ['CPUExecutionProvider']\n",
      "                                enable_ipu_cnn --- True\n",
      "                        enable_ipu_transformer --- False\n",
      "                     specific_tensor_precision --- False\n",
      "                                    debug_mode --- False\n",
      "                          convert_fp16_to_fp32 --- False\n",
      "                          convert_nchw_to_nhwc --- False\n",
      "                                   include_cle --- False\n",
      "                                    include_sq --- False\n",
      "                               include_fast_ft --- False\n",
      "                                 extra_options --- {'ActivationSymmetric': True}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:vai_q_onnx.quant_utils:The input ONNX model C:/Users/vgods/AppData/Local/Temp/vai.simp.kpf9kmm3/model_simp.onnx can run inference successfully\n",
      "INFO:vai_q_onnx.quantize:optimize the model for better hardware compatibility.\n",
      "INFO:vai_q_onnx.quantize:Start calibration...\n",
      "INFO:vai_q_onnx.quantize:Start collecting data, runtime depends on your model size and the number of calibration dataset.\n",
      "INFO:vai_q_onnx.calibrate:Finding optimal threshold for each tensor using PowerOfTwoMethod.MinMSE algorithm ...\n",
      "INFO:vai_q_onnx.calibrate:Use all calibration data to calculate min mse\n",
      "Computing range: 100%|██████████| 10/10 [00:04<00:00,  2.30tensor/s]\n",
      "INFO:vai_q_onnx.quantize:Finished the calibration of PowerOfTwoMethod.MinMSE which costs 4.6s\n",
      "INFO:vai_q_onnx.qdq_quantizer:Remove QuantizeLinear & DequantizeLinear on certain operations(such as conv-relu).\n",
      "INFO:vai_q_onnx.refine:Adjust the quantize info to meet the compiler constraints\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Op Type              </span>┃<span style=\"font-weight: bold\"> Float Model                      </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ Conv                 │<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\"> 4                                </span>│\n",
       "│ Relu                 │<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\"> 4                                </span>│\n",
       "│ Constant             │<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\"> 1                                </span>│\n",
       "│ Add                  │<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\"> 1                                </span>│\n",
       "├──────────────────────┼──────────────────────────────────┤\n",
       "│ Quantized model path │<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\"> models/helloworld_quantized.onnx </span>│\n",
       "└──────────────────────┴──────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mOp Type             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mFloat Model                     \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ Conv                 │\u001b[1;38;5;46m \u001b[0m\u001b[1;38;5;46m4                               \u001b[0m\u001b[1;38;5;46m \u001b[0m│\n",
       "│ Relu                 │\u001b[1;38;5;46m \u001b[0m\u001b[1;38;5;46m4                               \u001b[0m\u001b[1;38;5;46m \u001b[0m│\n",
       "│ Constant             │\u001b[1;38;5;46m \u001b[0m\u001b[1;38;5;46m1                               \u001b[0m\u001b[1;38;5;46m \u001b[0m│\n",
       "│ Add                  │\u001b[1;38;5;46m \u001b[0m\u001b[1;38;5;46m1                               \u001b[0m\u001b[1;38;5;46m \u001b[0m│\n",
       "├──────────────────────┼──────────────────────────────────┤\n",
       "│ Quantized model path │\u001b[1;38;5;46m \u001b[0m\u001b[1;38;5;46mmodels/helloworld_quantized.onnx\u001b[0m\u001b[1;38;5;46m \u001b[0m│\n",
       "└──────────────────────┴──────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calibrated and quantized model saved at: models/helloworld_quantized.onnx\n"
     ]
    }
   ],
   "source": [
    "# `input_model_path` is the path to the original, unquantized ONNX model.\n",
    "input_model_path = \"models/helloworld.onnx\"\n",
    "\n",
    "# `output_model_path` is the path where the quantized model will be saved.\n",
    "output_model_path = \"models/helloworld_quantized.onnx\"\n",
    "\n",
    "vai_q_onnx.quantize_static(\n",
    "    input_model_path,\n",
    "    output_model_path,\n",
    "    calibration_data_reader=None,\n",
    "    quant_format=vai_q_onnx.QuantFormat.QDQ,\n",
    "    calibrate_method=vai_q_onnx.PowerOfTwoMethod.MinMSE,\n",
    "    activation_type=vai_q_onnx.QuantType.QUInt8,\n",
    "    weight_type=vai_q_onnx.QuantType.QInt8,\n",
    "    enable_ipu_cnn=True,\n",
    "    extra_options={'ActivationSymmetric': True}\n",
    ")\n",
    "\n",
    "print('Calibrated and quantized model saved at:', output_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Run Model\n",
    "\n",
    "#### CPU Run\n",
    "\n",
    "Before runnning the model on the NPU, we'll run the model on the CPU and get the execution time for comparison with the NPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the path to the quantized ONNZ Model\n",
    "quantized_model_path = r'./models/helloworld_quantized.onnx'\n",
    "model = onnx.load(quantized_model_path)\n",
    "\n",
    "# Create some random input data for testing\n",
    "input_data = np.random.uniform(low=-1, high=1, size=(batch_size, input_channels, input_size, input_size)).astype(np.float32)\n",
    "\n",
    "cpu_options = onnxruntime.SessionOptions()\n",
    "\n",
    "# Create Inference Session to run the quantized model on the CPU\n",
    "cpu_session = onnxruntime.InferenceSession(\n",
    "    model.SerializeToString(),\n",
    "    providers = ['CPUExecutionProvider'],\n",
    "    sess_options=cpu_options,\n",
    ")\n",
    "\n",
    "# Run Inference\n",
    "start = timer()\n",
    "cpu_results = cpu_session.run(None, {'input': input_data})\n",
    "cpu_total = timer() - start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### NPU Run\n",
    "\n",
    "Now, we'll run it on the NPU and time the execution so that we can compare the results with the CPU.\n",
    "If the model has already been compiled, it won't recompile unless you delete the generated cache folder using the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Directory deleted successfully. Starting Fresh.\n"
     ]
    }
   ],
   "source": [
    "# We want to make sure we compile everytime, otherwise the tools will use the cached version\n",
    "# Get the current working directory\n",
    "current_directory = os.getcwd()\n",
    "directory_path = os.path.join(current_directory,  r'cache\\hello_cache')\n",
    "cache_directory = os.path.join(current_directory,  r'cache')\n",
    "\n",
    "# Check if the directory exists and delete it if it does.\n",
    "if os.path.exists(directory_path):\n",
    "    shutil.rmtree(directory_path)\n",
    "    print(f\"Directory deleted successfully. Starting Fresh.\")\n",
    "else:\n",
    "    print(f\"Directory '{directory_path}' does not exist.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compile and run\n",
    "\n",
    "On the first run, the model will compile for the NPU before executing the inference. It's best to run the following cell again if you want to see better inference times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "install_dir = os.environ['RYZEN_AI_INSTALLATION_PATH']\n",
    "config_file_path = os.path.join(install_dir, 'voe-4.0-win_amd64', 'vaip_config.json') # Path to the NPU config file\n",
    "\n",
    "aie_options = onnxruntime.SessionOptions()\n",
    "\n",
    "aie_session = onnxruntime.InferenceSession(\n",
    "    model.SerializeToString(),\n",
    "    providers=['VitisAIExecutionProvider'],\n",
    "    sess_options=aie_options,\n",
    "    provider_options = [{'config_file': config_file_path,\n",
    "                         'cacheDir': cache_directory,\n",
    "                         'cacheKey': 'hello_cache'}]\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Inference\n",
    "start = timer()\n",
    "npu_results = aie_session.run(None, {'input': input_data})\n",
    "npu_total = timer() - start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's gather our results and see what we have"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU Execution Time: 0.11257850000004055\n",
      "NPU Execution Time: 0.08555689999997185\n"
     ]
    }
   ],
   "source": [
    "print(f\"CPU Execution Time: {cpu_total}\")\n",
    "print(f\"NPU Execution Time: {npu_total}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** For a model this small in size, you likely won't see much of a performance gain when using the NPU versus the CPU. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at running the model on the NPU lots of times so that we can see the NPU being utilized.\n",
    "To do this, make sure to have Task Manager opened in a window you can see when you run the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterations = 50 # edit this for more or less\n",
    "\n",
    "for i in range(iterations):\n",
    "    npu_results = aie_session.run(None, {'input': input_data})\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And there you have it. Your first model running on the NPU. We recommend trying a more complex model like ResNet50 or a custom model to compare performance and accuracy on the NPU.\n"
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}