{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "license-cell",
   "metadata": {},
   "source": [
    "```\n",
    "SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
    "SPDX-License-Identifier: Apache-2.0\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\"); \n",
    "you may not use this file except in compliance with the License. \n",
    "You may obtain a copy of the License at\n",
    "\n",
    "http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbb82290-07cd-4b09-afa4-cf4ad3002f06",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# ResNet50 with Weight Stripping\n",
    "\n",
    "NVIDIA® [TensorRT](https://developer.nvidia.com/tensorrt)™ is an SDK designed for high-performance deep learning inference. TensorRT acts as an Ahead-of-Time (AoT) compiler, which necessitates the offline generation of an engine plan file. This file contains optimized backend kernels tailored for the production environment, with both the kernel code and weights embedded within it. \n",
    "\n",
    "For Independent Software Vendors (ISVs) deploying these engine plan files, it is crucial to have one engine file per SKU to achieve optimal performance. When there are N SKUs, the weights are duplicated N times, significantly increasing the installer size. This is particularly problematic for the Windows platform, which has over 40 consumer SKUs. \n",
    "\n",
    "To address this issue, we introduced the weight stripping feature in TensorRT 10.0. With weight stripping enabled, TensorRT creates a refittable engine based on the assumption that the refit weights will be identical to those provided during the build. The resulting engine maintains the same performance as a non-refittable one. All refittable weights can be adjusted through the refit API. Moreover, by stripping all the refittable weights, we produce a small plan file, with weights supplied later via refitting. This approach allows the use of a single set of weights across different inference backends or TensorRT plans for multiple GPU architectures. \n",
    "\n",
    "This notebook demonstrates how to use TensorRT to build a weight-stripped engine and later refit it to a full engine for inference, using ResNet-50 as an example.\n",
    "\n",
    "## Content\n",
    "1. [Download model and install pre-requisites](#1)\n",
    "2. [Building a full TensorRT engine](#2)\n",
    "3. [Building a weight-stripped TensorRT engine](#3)\n",
    "4. [Running inference example using the full engine](#4)\n",
    "5. [Running inference example using the weight-stripped engine](#5)\n",
    "6. [Serialize the refitted weight-stripped engine (optional)](#6)\n",
    "7. [Inference benchmarking](#7)\n",
    "8. [Engine plan size comparison](#8)\n",
    "9. [Advanced topic: weight-stripping with lean runtime](#9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae7cd6b9-0bd0-4e3c-9fbc-fc8c76201ffe",
   "metadata": {},
   "source": [
    "<a id=\"1\"></a>\n",
    "\n",
    "## 1. Download model and install pre-requisites\n",
    "\n",
    "### Download model and dataset\n",
    "\n",
    "First, let's download the ResNet50 ONNX model and test image from TensorRT installer package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6e632ed2-ee13-48aa-80f0-16542a0fe798",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-05-20 20:02:41--  https://download.onnxruntime.ai/onnx/models/resnet50.tar.gz\n",
      "Resolving download.onnxruntime.ai (download.onnxruntime.ai)... 13.107.246.69, 13.107.213.69, 2620:1ec:46::69, ...\n",
      "Connecting to download.onnxruntime.ai (download.onnxruntime.ai)|13.107.246.69|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 101632129 (97M) [application/octet-stream]\n",
      "Saving to: ‘resnet50.tar.gz’\n",
      "\n",
      "resnet50.tar.gz     100%[===================>]  96.92M  7.35MB/s    in 14s     \n",
      "\n",
      "2024-05-20 20:02:56 (6.87 MB/s) - ‘resnet50.tar.gz’ saved [101632129/101632129]\n",
      "\n",
      "--2024-05-20 20:03:00--  https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.0.1/tars/TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz\n",
      "Resolving developer.nvidia.com (developer.nvidia.com)... 152.195.19.142\n",
      "Connecting to developer.nvidia.com (developer.nvidia.com)|152.195.19.142|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://developer.download.nvidia.com/compute/machine-learning/tensorrt/10.0.1/tars/TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz [following]\n",
      "--2024-05-20 20:03:00--  https://developer.download.nvidia.com/compute/machine-learning/tensorrt/10.0.1/tars/TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz\n",
      "Resolving developer.download.nvidia.com (developer.download.nvidia.com)... 152.195.19.142\n",
      "Connecting to developer.download.nvidia.com (developer.download.nvidia.com)|152.195.19.142|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 2349001911 (2.2G) [application/x-gzip]\n",
      "Saving to: ‘TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz’\n",
      "\n",
      "TensorRT-10.0.1.6.L 100%[===================>]   2.19G   227MB/s    in 10s     \n",
      "\n",
      "2024-05-20 20:03:12 (215 MB/s) - ‘TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz’ saved [2349001911/2349001911]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://download.onnxruntime.ai/onnx/models/resnet50.tar.gz -O resnet50.tar.gz\n",
    "!tar -xzf resnet50.tar.gz\n",
    "\n",
    "!wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.0.1/tars/TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz\n",
    "!tar -xzf TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz TensorRT-10.0.1.6/data/int8_api/reference_labels.txt TensorRT-10.0.1.6/data/int8_api/airliner.ppm\n",
    "!mv TensorRT-10.0.1.6/data/int8_api/airliner.ppm  .\n",
    "!mv TensorRT-10.0.1.6/data/int8_api/reference_labels.txt  .\n",
    "!rm -rf TensorRT-10.0.1.6 TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932a8c50-61ef-4914-9941-05ad411b6527",
   "metadata": {},
   "source": [
    "### Install pre-requisites\n",
    "\n",
    "The Python sample is built on top of the native TensorRT API. To streamline the code and enhance its readability, we use polygraphy in the notebook. We also use matpoltlib to visualize the data. Therefore, it's necessary to install packages from both the requirement.txt file, polygraphy and matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "66a9caf4-aabc-42bd-a387-95a7eba342a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -q -r ../requirements.txt;\n",
    "%pip install -q matplotlib colored polygraphy --extra-index-url https://pypi.ngc.nvidia.com;\n",
    "import sys\n",
    "import site\n",
    "sys.path.append(site.getusersitepackages())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b22f43cb-d239-4323-ab10-b83b5f3c2095",
   "metadata": {},
   "source": [
    "<a id=\"2\"></a>\n",
    "\n",
    "## 2. Building a full TensorRT engine\n",
    "\n",
    "In this section, we will optimize the ResNet model for inference with strongly-typed enabled. The final engine will be saved to the local disk for subsequent engine plan size comparisons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b888e4-fb2e-45a2-bf03-452a976b52aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Configuring with profiles:[\n",
      "        Profile 0:\n",
      "            {gpu_0/data_0 [min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224]]}\n",
      "    ]\n",
      "\u001b[38;5;14m[I] Building engine with configuration:\n",
      "    Flags                  | []\n",
      "    Engine Capability      | EngineCapability.STANDARD\n",
      "    Memory Pools           | [WORKSPACE: 40326.38 MiB, TACTIC_DRAM: 40326.38 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]\n",
      "    Tactic Sources         | [EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]\n",
      "    Profiling Verbosity    | ProfilingVerbosity.DETAILED\n",
      "    Preview Features       | []\u001b[0m\n",
      "\u001b[38;5;10m[I] Finished engine building in 38.595 seconds\u001b[0m\n",
      "[I] Saving engine to resnet50_full.plan\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorrt.tensorrt.ICudaEngine at 0x7f41ed336170>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from polygraphy.backend.trt import CreateConfig, network_from_onnx_path, engine_from_network, save_engine\n",
    "\n",
    "# Create TensorRT engine from onnx model.\n",
    "network = network_from_onnx_path('resnet50/model.onnx', strongly_typed=True)\n",
    "config = CreateConfig()\n",
    "engine = engine_from_network(network, config)\n",
    "save_engine(engine, path='resnet50_full.plan')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c608c03-c835-4778-aea4-cbe1f94c74ef",
   "metadata": {},
   "source": [
    "<a id=\"3\"></a>\n",
    "\n",
    "## 3. Building a weight-stripped TensorRT engine\n",
    "\n",
    "To enable weight stripping, simply set the strip_plan flag in the builder configuration. This requires only one line of code modification as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19483305-5006-4f00-a0d3-5de6d6504406",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Configuring with profiles:[\n",
      "        Profile 0:\n",
      "            {gpu_0/data_0 [min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224]]}\n",
      "    ]\n",
      "\u001b[38;5;14m[I] Building engine with configuration:\n",
      "    Flags                  | [STRIP_PLAN]\n",
      "    Engine Capability      | EngineCapability.STANDARD\n",
      "    Memory Pools           | [WORKSPACE: 40326.38 MiB, TACTIC_DRAM: 40326.38 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]\n",
      "    Tactic Sources         | [EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]\n",
      "    Profiling Verbosity    | ProfilingVerbosity.DETAILED\n",
      "    Preview Features       | []\u001b[0m\n",
      "\u001b[38;5;10m[I] Finished engine building in 34.624 seconds\u001b[0m\n",
      "[I] Saving engine to resnet50_stripped.plan\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorrt.tensorrt.ICudaEngine at 0x7f41ed3599b0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create TensorRT engine from onnx model with weight-stripped.\n",
    "config = CreateConfig(strip_plan=True)\n",
    "engine = engine_from_network(network, config)\n",
    "save_engine(engine, path='resnet50_stripped.plan')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af4d7e3c-1f9d-42c9-ba3c-8121a0c085f0",
   "metadata": {},
   "source": [
    "<a id=\"4\"></a>\n",
    "\n",
    "## 4. Running inference example using the full engine\n",
    "\n",
    "In this section, we will use an airliner image to verify the inference results on both engines.\n",
    "\n",
    "### Prepare input data\n",
    "Let's use pillow to load the image and load the reference labels from a text file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f1d42483-1176-407d-89ca-d45506f12db5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 8588 characters]",
      "image/png": "[encoded data, 96200 characters]",
      "text/plain": [
       "<PIL.PpmImagePlugin.PpmImageFile image mode=RGB size=224x224>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "from IPython.display import display\n",
    "\n",
    "# Load the image.\n",
    "image_path = 'airliner.ppm'\n",
    "img = Image.open(image_path)\n",
    "\n",
    "# Display the image.\n",
    "display(img)\n",
    "\n",
    "# Transpose to NCHW layout and normalize data.\n",
    "img = np.transpose(img, (2, 0, 1))  \n",
    "img = np.array(img).astype(np.float32) / 255.0\n",
    "img = np.expand_dims(img, axis=0)\n",
    "\n",
    "# Load labels\n",
    "labels = open('reference_labels.txt', \"r\").read().split(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef7b47f6-6f2b-4edf-b997-b62df1ec36cc",
   "metadata": {},
   "source": [
    "### Load full TensorRT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "776fd81f-d309-4f73-8fde-57d7928bc929",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Loading bytes from resnet50_full.plan\n"
     ]
    }
   ],
   "source": [
    "import tensorrt as trt\n",
    "from polygraphy.backend.common import bytes_from_path\n",
    "from polygraphy.backend.trt import engine_from_bytes\n",
    "\n",
    "# Load TensorRT engine from plan file.\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.WARNING)\n",
    "full_engine = engine_from_bytes(bytes_from_path('resnet50_full.plan'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f78e141b-22be-415d-8dcb-ae3192102dda",
   "metadata": {},
   "source": [
    "### Inference with full engine\n",
    "Let's verify that TensorRT can recognize this image as 'airliner'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d4584157-09f1-45d0-89af-d6bb408b670d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full engine recognize image as 'airliner'!\n"
     ]
    }
   ],
   "source": [
    "from polygraphy.backend.trt import TrtRunner\n",
    "\n",
    "# Inference and get the result.\n",
    "with TrtRunner(full_engine) as runner:\n",
    "    outputs = runner.infer(feed_dict={\"gpu_0/data_0\": img})\n",
    "\n",
    "pred = labels[np.argmax(outputs['gpu_0/softmax_1'])]\n",
    "print(f\"Full engine recognize image as '{pred}'!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5d5027b-a68c-4103-beda-46ee6574e514",
   "metadata": {},
   "source": [
    "<a id=\"5\"></a>\n",
    "\n",
    "## 5. Running inference example using the weight-stripped engine\n",
    "\n",
    "Now, we will perform inference again using the weight-stripped engine. Let's skip the input data preparation step since it's the same as the full engine inference, and jump straight to loading the engine.\n",
    "\n",
    "### Load weight-stripped engine and update weights\n",
    "Note that this time, all the weights are stripped in the engine. We cannot directly use this engine for inference; instead, we need to pull the weights back using the refit API first. Since the engine was built using an ONNX model, we can use the OnnxParserRefitter class provided by TensorRT to automate the refit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "51cb8223-9127-4c17-9734-50cb793c3c5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Loading bytes from resnet50_stripped.plan\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load TensorRT engine from plan file.\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.WARNING)\n",
    "stripped_engine = engine_from_bytes(bytes_from_path('resnet50_stripped.plan'))\n",
    "\n",
    "# Update weights using refit API.\n",
    "refitter = trt.Refitter(stripped_engine, TRT_LOGGER)\n",
    "parser_refitter = trt.OnnxParserRefitter(refitter, TRT_LOGGER)\n",
    "parser_refitter.refit_from_file('resnet50/model.onnx')\n",
    "refitter.refit_cuda_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481303d9-d774-4b76-828c-8b170729ce1d",
   "metadata": {},
   "source": [
    "### Inference with refitted weight-stripped engine\n",
    "Let's verify that TensorRT can still recognize this image as 'airliner' on the refitted weight-stripped engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "42fe4c5f-108e-4d7f-ae80-a569c76dfe6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Refitted stripped engine recognize image as 'airliner'!\n"
     ]
    }
   ],
   "source": [
    "# Inference and get the result.\n",
    "with TrtRunner(stripped_engine) as runner:\n",
    "    outputs = runner.infer(feed_dict={\"gpu_0/data_0\": img})\n",
    "\n",
    "pred = labels[np.argmax(outputs['gpu_0/softmax_1'])]\n",
    "print(f\"Refitted stripped engine recognize image as '{pred}'!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7736d41-1703-4482-96f8-1d1b6ee24018",
   "metadata": {},
   "source": [
    "<a id=\"6\"></a>\n",
    "\n",
    "## 6. Serialize the refitted weight-stripped engine (optional) \n",
    "We can save the refitted, weight-stripped engine to disk, so you won't need to refit it next time you run inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6f11ebb7-346d-4d4e-963b-9339d8e4e639",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create config that include weights in plan.\n",
    "s_config = stripped_engine.create_serialization_config()\n",
    "s_config.flags &= ~(1 << int(trt.SerializationFlag.EXCLUDE_WEIGHTS))\n",
    "\n",
    "# Serialize the plan file.\n",
    "full_binary = engine.serialize_with_config(s_config)\n",
    "with open('resnet50_stripped_updated.plan', 'wb') as f:\n",
    "    f.write(full_binary)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "705d92bd-4f16-4c50-974d-b414f7314421",
   "metadata": {},
   "source": [
    "<a id=\"7\"></a>\n",
    "\n",
    "## 7. Inference benchmarking\n",
    "We have just finished verifying the accuracy. Now, let's measure the performance between the full engine and the weight-stripped engine.\n",
    "\n",
    "### Benchmarking a full engine\n",
    "First, let's check the total time of 100 iterations on the full engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "082e336c-0f8e-4d48-9d69-28dcd6fd9b0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full engine inference time on 100 iterations: 0.2675 seconds\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "# Inference for 100 iterations.\n",
    "start_time = time.time()\n",
    "with TrtRunner(full_engine) as runner:\n",
    "    for i in range(100):\n",
    "        outputs = runner.infer(feed_dict={\"gpu_0/data_0\": img})\n",
    "total_time_full = time.time() - start_time\n",
    "print(\"Full engine inference time on 100 iterations: {:.4f} seconds\".format(total_time_full))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dbbaf7b-9a70-4478-b675-13656f9b272a",
   "metadata": {},
   "source": [
    "### Benchmarking a weight-stripped engine\n",
    "Again, let's use the same code to check the 100 iterations on the weight-stripped engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b48f4c09-ab6a-4a6b-bf36-628edf2ca7de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Refitted stripped engine inference time on 100 iterations: 0.2671 seconds\n"
     ]
    }
   ],
   "source": [
    "# Inference for 100 iterations.\n",
    "start_time = time.time()\n",
    "with TrtRunner(stripped_engine) as runner:\n",
    "    for i in range(100):\n",
    "        outputs = runner.infer(feed_dict={\"gpu_0/data_0\": img})\n",
    "total_time_stripped = time.time() - start_time\n",
    "print(\"Refitted stripped engine inference time on 100 iterations: {:.4f} seconds\".format(total_time_stripped))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dc64563-0088-4682-ba97-8c84523ca1bd",
   "metadata": {},
   "source": [
    "Let's do a comparison, but let's ignore the small variances due to using the floating clock. We can see that there is no performance difference between the full engine and the weight-stripped engine!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f4f8c300-eb3a-4b59-958c-5332e8dc546f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27196 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Plot the inference time comparison.\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.bar(['full', 'weight-stripped'], [total_time_full, total_time_stripped], color=['blue', 'green'])\n",
    "plt.xlabel('Build configurations')\n",
    "plt.ylabel('Time / 100 iterations (sec)')\n",
    "plt.title('Performance Comparison')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec6b666c-2eee-43a0-a6f5-435083ecadf7",
   "metadata": {},
   "source": [
    "<a id=\"8\"></a>\n",
    "\n",
    "## 8. Engine plan size comparison\n",
    "So far, we have verified that we did not lose performance or accuracy for the weight-stripped engine. Now, let's check the engine file size. With simple math, we can see that we reduced 95% of the original engine size. Considering deployment on N SKUs, we will save 95% on each SKU. This will greatly help with the TensorRT integration into your own project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ac44f6df-0983-45ec-a33f-4a30a2df1957",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22364 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os \n",
    "\n",
    "# Get the file sizes in MiB\n",
    "full_size = os.path.getsize('resnet50_full.plan') / 1024 / 1024\n",
    "stripped_size = os.path.getsize('resnet50_stripped.plan') / 1024 / 1024\n",
    "\n",
    "# Plot the engine size comparison.\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.bar(['full', 'weight-stripped'], [full_size, stripped_size], color=['blue', 'green'])\n",
    "plt.xlabel('Build configurations')\n",
    "plt.ylabel('Size (MiB)')\n",
    "plt.title('Engine Size Comparison')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63849d12-4565-457e-a64c-724935f3f276",
   "metadata": {},
   "source": [
    "<a id=\"9\"></a>\n",
    "\n",
    "## 9. Advanced topic: weight-stripping with lean runtime\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7498ff-d879-48e5-ab00-b7c2c3c1da85",
   "metadata": {},
   "source": [
    "Additionally, we can leverage the lean runtime to further reduce the package size for the weight-stripped engine. The lean runtime is the same runtime used in version-compatible engines. Originally designed to enable users to generate a TensorRT engine with version X and load it with an application built with version Y, the lean runtime library is relatively small, approximately 40 MiB. Therefore, software distributors on top of TensorRT only need to ship the weight-stripped engine along with the 40 MiB lean runtime when the weights are already available on the target customer machine.\n",
    "\n",
    "### Build with version compatible\n",
    "This can be achieved with as little as one line change to the builder code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79b7b4be-1377-48e5-9a45-ad249fec04af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Configuring with profiles:[\n",
      "        Profile 0:\n",
      "            {gpu_0/data_0 [min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224]]}\n",
      "    ]\n",
      "[I] Version or hardware compatibility was enabled. If you are using an ONNX model, please set the NATIVE_INSTANCENORM ONNX parser flag, e.g. `--onnx-flags NATIVE_INSTANCENORM`\n",
      "\u001b[38;5;14m[I] Building engine with configuration:\n",
      "    Flags                  | [VERSION_COMPATIBLE, EXCLUDE_LEAN_RUNTIME]\n",
      "    Engine Capability      | EngineCapability.STANDARD\n",
      "    Memory Pools           | [WORKSPACE: 40326.38 MiB, TACTIC_DRAM: 40326.38 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]\n",
      "    Tactic Sources         | [EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]\n",
      "    Profiling Verbosity    | ProfilingVerbosity.DETAILED\n",
      "    Preview Features       | []\u001b[0m\n",
      "\u001b[38;5;10m[I] Finished engine building in 38.043 seconds\u001b[0m\n",
      "[I] Saving engine to resnet50_full_vc.plan\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorrt.tensorrt.ICudaEngine at 0x7f40a4359970>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build TensorRT full engine with version compatible.\n",
    "config = CreateConfig(version_compatible=True, exclude_lean_runtime=True)\n",
    "engine = engine_from_network(network, config)\n",
    "save_engine(engine, path='resnet50_full_vc.plan')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb64a682-148e-4599-b0ce-24ff2094831c",
   "metadata": {},
   "source": [
    "### Build with both weight-stripped and version compatible\n",
    "This can also be accomplished with just a single line change to the builder code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9fc753b-fc1e-470d-96a7-f7fafaabca0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Configuring with profiles:[\n",
      "        Profile 0:\n",
      "            {gpu_0/data_0 [min=[1, 3, 224, 224], opt=[1, 3, 224, 224], max=[1, 3, 224, 224]]}\n",
      "    ]\n",
      "[I] Version or hardware compatibility was enabled. If you are using an ONNX model, please set the NATIVE_INSTANCENORM ONNX parser flag, e.g. `--onnx-flags NATIVE_INSTANCENORM`\n",
      "\u001b[38;5;14m[I] Building engine with configuration:\n",
      "    Flags                  | [VERSION_COMPATIBLE, EXCLUDE_LEAN_RUNTIME, STRIP_PLAN]\n",
      "    Engine Capability      | EngineCapability.STANDARD\n",
      "    Memory Pools           | [WORKSPACE: 40326.38 MiB, TACTIC_DRAM: 40326.38 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]\n",
      "    Tactic Sources         | [EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS]\n",
      "    Profiling Verbosity    | ProfilingVerbosity.DETAILED\n",
      "    Preview Features       | []\u001b[0m\n",
      "\u001b[38;5;10m[I] Finished engine building in 36.902 seconds\u001b[0m\n",
      "[I] Saving engine to resnet50_stripped_vc.plan\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorrt.tensorrt.ICudaEngine at 0x7f3fe0f5ca70>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build TensorRT weight-stripped engine with version compatible.\n",
    "config = CreateConfig(strip_plan=True, version_compatible=True, exclude_lean_runtime=True)\n",
    "engine = engine_from_network(network, config)\n",
    "save_engine(engine, path='resnet50_stripped_vc.plan')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8d1e6c-dd81-41bb-8c51-e3e2a872e597",
   "metadata": {},
   "source": [
    "### Load weight-stripped engine with lean runtime\n",
    "\n",
    "The dispatch runtime and lean runtime are the dependency libraries you need to ship with your software. Here, the dispatch runtime serves as a shim library, forwarding API calls to the lean runtime. Meanwhile, the lean runtime contains both the host code and device code necessary to launch your inference job.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cc9a7dd1-f614-413d-8cf2-745615ee16d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Loading TensorRT runtime from: /usr/lib/x86_64-linux-gnu/libnvinfer_lean.so\n",
      "[I] Loading bytes from resnet50_stripped_vc.plan\n"
     ]
    }
   ],
   "source": [
    "from polygraphy.backend.trt import load_runtime\n",
    "\n",
    "# Load the weight-stripped engine using lean runtime.\n",
    "lean_runtime = load_runtime('/usr/lib/x86_64-linux-gnu/libnvinfer_lean.so')\n",
    "stripped_vc_engine = engine_from_bytes(bytes_from_path('resnet50_stripped_vc.plan'), runtime=lean_runtime)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2e5a7fd-e921-4ad6-b65a-3d6f38eab4ad",
   "metadata": {},
   "source": [
    "### Update weights using refit API\n",
    "\n",
    "We still use the same refit API to update the weights loaded from the lean runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "888861e6-03d2-4997-8023-d4096b37e545",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Update the weight-stripped engine using refit API.\n",
    "refitter = trt.Refitter(stripped_vc_engine, TRT_LOGGER)\n",
    "parser_refitter = trt.OnnxParserRefitter(refitter, TRT_LOGGER)\n",
    "parser_refitter.refit_from_file('resnet50/model.onnx')\n",
    "refitter.refit_cuda_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c64e4686-41c4-451b-b90e-905246372f19",
   "metadata": {},
   "source": [
    "### Inference with airliner image\n",
    "Now it is time to verify the result!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "bc9ee086-cc92-464c-b46a-7d379265c174",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Refitted stripped VC engine recognize image as 'airliner'!\n"
     ]
    }
   ],
   "source": [
    "# Inference and get the result.\n",
    "with TrtRunner(stripped_vc_engine) as runner:\n",
    "    outputs = runner.infer(feed_dict={\"gpu_0/data_0\": img})\n",
    "\n",
    "pred = labels[np.argmax(outputs['gpu_0/softmax_1'])]\n",
    "print(f\"Refitted stripped VC engine recognize image as '{pred}'!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "733616c8-4d66-4987-9c40-6c1aec7fe751",
   "metadata": {},
   "source": [
    "### Size comparison of different deployment solution\n",
    "Finally let's compare the four solutions here:\n",
    "- Full engine + full runtime\n",
    "- Weight-stripped engine + full runtime\n",
    "- Full engine + lean runtime + dispatch runtime\n",
    "- Weight-stripped engine + lean runtime + dispatch runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "46a86192-e11c-481c-b98e-900aab72d3f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27896 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get the library sizes in MiB.\n",
    "full_runtime_size = os.path.getsize(os.path.realpath('/usr/lib/x86_64-linux-gnu/libnvinfer.so')) / 1024 / 1024\n",
    "lean_runtime_size = os.path.getsize(os.path.realpath('/usr/lib/x86_64-linux-gnu/libnvinfer_lean.so')) / 1024 / 1024\n",
    "dispatch_runtime_size = os.path.getsize(os.path.realpath('/usr/lib/x86_64-linux-gnu/libnvinfer_dispatch.so')) / 1024 / 1024\n",
    "\n",
    "# Get the library + engine file sizes in MiB.\n",
    "full_size = full_runtime_size + os.path.getsize('resnet50_full.plan') / 1024 / 1024\n",
    "stripped_size = full_runtime_size + os.path.getsize('resnet50_stripped.plan') / 1024 / 1024\n",
    "full_vc_size = lean_runtime_size + dispatch_runtime_size + os.path.getsize('resnet50_full_vc.plan') / 1024 / 1024\n",
    "stripped_vc_size = lean_runtime_size + dispatch_runtime_size + os.path.getsize('resnet50_stripped.plan') / 1024 / 1024\n",
    "\n",
    "# Plot the distribution size comparison.\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.bar(['full', 'weight-stripped', 'full + VC', 'weight-stripped + VC'], [full_size, stripped_size, full_vc_size, stripped_vc_size], color=['blue', 'green', 'purple', 'red'])\n",
    "plt.xlabel('Build configurations')\n",
    "plt.ylabel('Size (MiB)')\n",
    "plt.title('Distribute Package Size Comparison')\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
