07-model-optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b88dd1a1-c3ab-44ba-85db-2f72918f50f9",
   "metadata": {},
   "source": [
    "# Model Optimization Strategy"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "086116ed-31ff-4206-918e-a468e0ac4527",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Why We Need Optimize the Model\n",
    "Model optimization methods are **crucial for reducing latency**, **achieving higher throughput**, and **maintaining performance metrics** in production environments. However, they also come with their own challenges.\n",
    "\n",
    "Take a closer look at the animation below. It's a short video clip made up of different *Frame Per Second* (**FPS**) rates. When the FPS rate is higher, the video has clearer and more detailed features. On the other hand, lower FPS rates result in a lagging video that isn't enjoyable to watch and lacks important details for the application.\n",
    "\n",
    "Imagine if we created an application that counts the number of vehicles or people passing through a gate. If the FPS rate is low, we would miss many objects because the video stream is slow. Similarly, if we want to detect burglars to trigger alarms, our application may not be able to capture them in time because they appear blurry with the slower FPS rate. Having a higher FPS rate is extremely important for certain applications.\n",
    "\n",
    "<center><img src=\"https://i.makeagif.com/media/10-05-2015/1VSTG7.gif\" width=600></center>\n",
    "\n",
    "Ideally, we aim to process the video in real-time at around 30 frames per second (FPS) from start to finish. However, in reality, there are multiple steps involved in this process. First, we need to decode the image, then process it using the model, apply other postprocessing techniques, and display the final result. Additionally, there is a step that may not be crucial but still requires time and processing power, which involves copying or moving data between the central processing unit (CPU) and the graphics processing unit (GPU).\n",
    "\n",
    "<center><img src=\"figures/cuda-flow.png\" width=600></center>\n",
    "\n",
    "When we want to maximize computation **using the GPU**, we typically need to *transfer the data from the CPU to the GPU and vice versa*. This back-and-forth movement of data between the CPU and GPU **can be time-consuming and unnecessary**. What if we could process everything before using the final result on the CPU? By ensuring that the data is directly moved to the GPU, preparing and processing the model on the GPU, and finally returning the result to the CPU, we can avoid multiple unnecessary data transfers between the CPU and GPU.\n",
    "\n",
    "The model inference or computation is the most expensive load, that is why we're focusing on optimizing the model first. After that, we can start optimizing other processes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4794e73b-1feb-4213-a99e-d9a10463d910",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Optimization Strategy\n",
    "\n",
    "We are primarily concerned with the usual ways to make the model better while it's being trained and after it's done training. Model compression techniques can be used throughout and after training the model, while parallelization and hardware acceleration can be employed after the training is completed.\n",
    "\n",
    "### Model Compression\n",
    "\n",
    "Model compression techniques are methods used to **reduce the size and computational complexity of machine learning models** while maintaining their performance as much as possible. These techniques are particularly useful in scenarios where memory and computational resources are limited, such as deploying models on edge devices or mobile applications. It's important to note that while model compression techniques reduce model size and computational requirements, there might be a trade-off with performance metrics such as accuracy or inference speed. Here are some commonly used model compression techniques:\n",
    "\n",
    "- **Pruning**: Pruning involves removing unnecessary connections or weights from the model, reducing its size and computational requirements. There are two main types of pruning: weight pruning and structured pruning. Weight pruning involves removing individual weights below a certain threshold, while structured pruning removes entire neurons, channels, or layers. Pruning can be performed during or after training, and it can achieve significant model size reduction with minimal impact on performance.\n",
    "\n",
    "- **Quantization**: Quantization reduces the precision of numerical values in the model, such as weights or activations. By representing numbers with fewer bits, the model's size is reduced, leading to reduced memory footprint and faster computations. For example, converting 32-bit floating-point values to 8-bit integers can achieve a 4x reduction in size. Quantization can introduce some performance degradation, but techniques like post-training quantization or quantization-aware training aim to minimize this impact.\n",
    "\n",
    "- **Knowledge Distillation**: Knowledge distillation involves training a smaller, more lightweight model (student model) to mimic the behavior and predictions of a larger, more complex model (teacher model). The student model learns from the teacher model's outputs, using them as soft targets during training. This technique allows transferring the knowledge captured by the larger model to the smaller one, resulting in a compact model with similar performance.\n",
    "\n",
    "- **Low-Rank Approximation**: Low-rank approximation aims to reduce the computational complexity of a model by approximating its weight matrices with lower-rank matrices. This technique leverages the fact that many weight matrices in neural networks are of high rank but can be well approximated by lower-rank matrices. By reducing the rank, the model's size and computational requirements are decreased, leading to faster inference.\n",
    "\n",
    "- **Factorization**: Factorization methods decompose weight matrices into two or more lower-dimensional matrices. Common factorization techniques include matrix factorization, tensor factorization, or singular value decomposition (SVD). Factorization can reduce the number of parameters and improve the model's efficiency while maintaining performance.\n",
    "\n",
    "- **Compact Architectures**: Designing compact architectures from scratch is another approach to model compression. These architectures are specifically designed to be lightweight and efficient while achieving good performance. Examples of compact architectures include MobileNet, ShuffleNet, or SqueezeNet. They often utilize techniques like depth-wise separable convolutions, channel shuffling, or bottleneck structures to reduce computational complexity.\n",
    "\n",
    "### Parallelization and Hardware Acceleration\n",
    "\n",
    "Parallelization and hardware acceleration techniques are employed to improve the efficiency and speed of machine learning models by leveraging the power of multiple processing units and specialized hardware. These techniques help in achieving higher throughput and reducing the overall latency of model inference. Here's an explanation of parallelization and hardware acceleration:\n",
    "\n",
    "- **Parallelization**: Parallelization involves dividing the computational workload of a machine learning model across multiple processing units, enabling them to work simultaneously. This can be achieved through two main approaches:\n",
    "\n",
    "  - **Data Parallelism**: Data parallelism involves replicating the model across multiple processing units and dividing the training or inference data among them. Each processing unit independently performs computations on its assigned data subset. The results are then combined to obtain the final prediction. This approach is particularly useful for scenarios where the model is trained or evaluated on large datasets.\n",
    "\n",
    "  - **Model Parallelism**: Model parallelism is applied when the model's architecture is too large to fit within a single processing unit's memory. In this approach, different parts of the model are assigned to different processing units, and computations are performed in a coordinated manner across these units. This allows for the parallel execution of model computations while efficiently utilizing available resources.\n",
    "\n",
    "- **Hardware Acceleration**: Hardware acceleration involves utilizing specialized hardware components to expedite the computations required by machine learning models. Some common hardware acceleration techniques include:\n",
    "\n",
    "  - **Graphics Processing Units (GPUs)**: GPUs are highly parallel processors capable of executing multiple tasks simultaneously. They excel at performing matrix operations, which are fundamental to many machine learning algorithms. GPUs can significantly speed up training and inference by parallelizing computations across thousands of cores, enabling faster processing of large datasets.\n",
    "\n",
    "  - **Tensor Processing Units (TPUs)**: TPUs are Google's custom-designed application-specific integrated circuits (ASICs) optimized for machine learning workloads. They provide enhanced performance and energy efficiency compared to general-purpose CPUs and GPUs. TPUs are particularly useful for deep learning tasks, offering faster model training and inference with lower power consumption."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83c94aa6-3c36-40a9-8ee3-d2602b8efa7b",
   "metadata": {},
   "source": [
    "## Model Optimization Libraries\n",
    "\n",
    "Model optimization libraries are software tools that provide a range of techniques and functionalities to optimize deep learning models. These model optimization libraries and methods offer a range of techniques and tools to optimize deep learning models for improved performance, memory efficiency, accelerated inference, and deployment on specific hardware platforms. Here are some commonly used model optimization libraries:\n",
    "\n",
    "- **TorchScript**: TorchScript is a component of the PyTorch framework. It allows PyTorch models to be converted into an optimized and portable format suitable for deployment. TorchScript converts the model into a serialized representation that can be executed independently of the Python interpreter. It enables efficient execution of models, supports dynamic control flow, and provides tools for optimizing the model's performance, such as fusion of operations and support for quantization. TorchScript is particularly useful for deploying PyTorch models in production environments.\n",
    "\n",
    "- **TensorRT**: TensorRT is an inference optimizer and runtime library developed by NVIDIA. It is designed to accelerate deep learning models on NVIDIA GPUs. TensorRT optimizes models for high throughput and low latency by applying techniques like layer fusion, precision calibration, and dynamic tensor memory management. It also leverages GPU-specific optimizations to speed up inference. TensorRT supports various model formats, including ONNX, and provides an efficient runtime for accelerated inference on NVIDIA GPUs.\n",
    "\n",
    "- **OpenVino**: OpenVino (Open Visual Inference & Neural Network Optimization) is a toolkit provided by Intel. It enables optimization and deployment of deep learning models on Intel CPUs, GPUs, and FPGAs. OpenVino includes tools for model quantization, layer fusion, and hardware-specific optimizations. It provides an inference engine that maximizes the performance and efficiency of models on Intel hardware platforms. OpenVino supports popular deep learning frameworks and offers cross-platform deployment capabilities.\n",
    "\n",
    "- **ONNX (Open Neural Network Exchange)**: ONNX is an open format for representing machine learning models. It facilitates interoperability between different deep learning frameworks. With ONNX, models can be exported from one framework and imported into another for inference or further optimization. ONNX enables efficient exchange and deployment of models across different environments and platforms. It also provides tools like ONNX Runtime, an optimized engine for accelerated inference of ONNX models.\n",
    "\n",
    "- **DeepSpeed**: DeepSpeed is a library specifically designed for optimizing and accelerating training of deep learning models. It focuses on improving training speed, memory efficiency, and model scalability. DeepSpeed achieves this through various techniques such as memory optimization, gradient checkpointing, and offloading computation to specialized hardware like NVMe SSDs. It integrates with PyTorch and provides enhanced capabilities for training large-scale deep learning models.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b874a9df-edb1-4df1-9b0d-2ede5d5cc67b",
   "metadata": {},
   "source": [
    "## Comparing Different Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0322ddc7-28e5-4f7c-85a0-866ee5225777",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.1.2 is available.\n",
      "You should consider upgrading via the '/usr/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "! pip install -q timm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0293ea8-e2d7-4630-97d9-3fa2460dbe06",
   "metadata": {},
   "source": [
    "### Baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5f2a1f9c-9077-477f-8c9f-a57354d6e156",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on CPU (FP32): 0.12256813049316406\n",
      "Running on GPU (FP32): 0.02961587905883789\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import timm\n",
    "import torch\n",
    "\n",
    "model = timm.create_model(model_name='resnet50')\n",
    "model.eval()\n",
    "\n",
    "input = torch.rand(1, 3, 384, 384)\n",
    "\n",
    "def get_benchmark(model: torch.nn.Module, input: torch.Tensor, device: str):\n",
    "    device_ = torch.device(device)\n",
    "    model = model.to(device_)\n",
    "    input = input.to(device_)\n",
    "\n",
    "    # Warmup\n",
    "    for _ in range(10):\n",
    "        _ = model(input)\n",
    "\n",
    "    # Compute time\n",
    "    t = time.time()\n",
    "    _ = model(input)\n",
    "    latency = time.time() - t\n",
    "    return latency\n",
    "\n",
    "\n",
    "torch_cpu_time = 10\n",
    "torch_gpu_time = 10\n",
    "# CPU\n",
    "torch_cpu_time = get_benchmark(model, input, 'cpu')\n",
    "print(f'Running on CPU (FP32): {torch_cpu_time}')\n",
    "\n",
    "# GPU\n",
    "if torch.cuda.is_available():\n",
    "    torch_gpu_time = get_benchmark(model, input, 'cuda')\n",
    "    print(f'Running on GPU (FP32): {torch_gpu_time}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eedf940f-3f4e-4152-9b9e-5ba143fc0cca",
   "metadata": {},
   "source": [
    "### TorchScript Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "cb66a844-4b8a-4c01-996e-80d900580730",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on CPU (FP32): 0.1342177391052246 - Speedup: 0.913%\n",
      "Running on GPU (FP32): 0.02610039710998535 - Speedup: 1.135%\n"
     ]
    }
   ],
   "source": [
    "scripted_model = torch.jit.script(model)\n",
    "\n",
    "scripted_cpu_time = 10\n",
    "scripted_gpu_time = 10\n",
    "\n",
    "# CPU\n",
    "scripted_cpu_time = get_benchmark(scripted_model, input, 'cpu')\n",
    "print(f'Running on CPU (FP32): {scripted_cpu_time} - Speedup: {torch_cpu_time/scripted_cpu_time:.3f}%')\n",
    "\n",
    "# GPU\n",
    "if torch.cuda.is_available():\n",
    "    scripted_gpu_time = get_benchmark(scripted_model, input, 'cuda')\n",
    "    print(f'Running on GPU (FP32): {scripted_gpu_time} - Speedup: {torch_gpu_time/scripted_gpu_time:.3f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a427dfc5-ff2f-49e3-b6c4-4de13a8240d7",
   "metadata": {},
   "source": [
    "### ONNX Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "8b37b343-edc3-4a90-863d-894fd46c112f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.1.2 is available.\n",
      "You should consider upgrading via the '/usr/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -q onnxruntime onnxruntime-gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "b76f11c0-0e66-4826-b681-7cdd603320cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========== Diagnostic Run torch.onnx.export version 2.1.0a0+fe05266 ===========\n",
      "verbose: False, log level: Level.ERROR\n",
      "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
      "\n",
      "Running on CPU (FP32): 0.05630326271057129 - Speedup: 2.177%\n",
      "Running on CPU (FP32): 0.05630326271057129 - Speedup: 0.509%\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import onnx\n",
    "import onnxruntime\n",
    "import torch.onnx\n",
    "\n",
    "def to_numpy(tensor):\n",
    "    if isinstance(tensor, np.ndarray):\n",
    "        return tensor\n",
    "    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n",
    "\n",
    "def get_onnx_benchmark(model: str, input: np.array, device: str):\n",
    "    providers=['CPUExecutionProvider']\n",
    "    if 'cuda':\n",
    "        providers=['CUDAExecutionProvider', 'CPUExecutionProvider']\n",
    "    ort_session = onnxruntime.InferenceSession(model, providers=providers)\n",
    "\n",
    "    ort_inputs = {ort_session.get_inputs()[0].name: input}\n",
    "\n",
    "    # Warmup\n",
    "    for _ in range(10):\n",
    "        _ = ort_session.run(None, ort_inputs)\n",
    "\n",
    "    # Compute time\n",
    "    t = time.time()\n",
    "    _ = ort_session.run(None, ort_inputs)\n",
    "    latency = time.time() - t\n",
    "    return latency\n",
    "\n",
    "device_ = torch.device('cpu')\n",
    "model = model.to(device_)\n",
    "input = input.to(device_)\n",
    "\n",
    "# Export the model\n",
    "torch.onnx.export(model,               # model being run\n",
    "                  input,                         # model input (or a tuple for multiple inputs)\n",
    "                  \"model.onnx\",   # where to save the model (can be a file or file-like object)\n",
    "                  export_params=True,        # store the trained parameter weights inside the model file\n",
    "                  opset_version=15,          # the ONNX version to export the model to\n",
    "                  do_constant_folding=True,  # whether to execute constant folding for optimization\n",
    "                  input_names = ['input'],   # the model's input names\n",
    "                  output_names = ['output'], # the model's output names\n",
    "                 )\n",
    "\n",
    "onnx_model = onnx.load('model.onnx')\n",
    "onnx.checker.check_model(onnx_model)\n",
    "\n",
    "# CPU\n",
    "onnx_cpu_time = get_onnx_benchmark('model.onnx', to_numpy(input), 'cpu')\n",
    "print(f'Running on CPU (FP32): {onnx_cpu_time} - Speedup: {torch_cpu_time/onnx_cpu_time:.3f}%')\n",
    "\n",
    "onnx_gpu_time = get_onnx_benchmark('model.onnx', to_numpy(input), 'cuda')\n",
    "print(f'Running on CPU (FP32): {onnx_cpu_time} - Speedup: {torch_gpu_time/onnx_gpu_time:.3f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84c01393-59dc-4ce1-be7a-78e6819cfda5",
   "metadata": {},
   "source": [
    "### TensorRT Optimization\n",
    "\n",
    "TensorRT may not be available on all devices or systems, so you might have to set up your device with CUDA. Let me demonstrate the outcome of optimizing the same ONNX model with TensorRT, which can achieve a high frame rate of **60.291** frames per second (FPS) or process in just **0.016586223** seconds.\n",
    "\n",
    "```shell\n",
    "[06/14/2023-08:24:06] [I] === Device Information ===\n",
    "[06/14/2023-08:24:06] [I] Selected Device: NVIDIA GeForce GTX 1050\n",
    "[06/14/2023-08:24:06] [I] Compute Capability: 6.1\n",
    "[06/14/2023-08:24:06] [I] SMs: 5\n",
    "[06/14/2023-08:24:06] [I] Compute Clock Rate: 1.493 GHz\n",
    "[06/14/2023-08:24:06] [I] Device Global Memory: 4038 MiB\n",
    "[06/14/2023-08:24:06] [I] Shared Memory per SM: 96 KiB\n",
    "[06/14/2023-08:24:06] [I] Memory Bus Width: 128 bits (ECC disabled)\n",
    "[06/14/2023-08:24:06] [I] Memory Clock Rate: 3.504 GHz\n",
    "\n",
    "[06/14/2023-08:26:15] [I] === Trace details ===\n",
    "[06/14/2023-08:26:15] [I] Trace averages of 100 runs:\n",
    "[06/14/2023-08:26:15] [I] Average on 100 runs - GPU latency: 16.4322 ms - Host latency: 16.6029 ms (enqueue 1.60057 ms)\n",
    "[06/14/2023-08:26:15] [I] \n",
    "[06/14/2023-08:26:15] [I] === Performance summary ===\n",
    "[06/14/2023-08:26:15] [I] Throughput: 60.291 qps\n",
    "[06/14/2023-08:26:15] [I] Latency: min = 16.4145 ms, max = 17.2104 ms, mean = 16.6728 ms, median = 16.7601 ms, percentile(90%) = 16.8579 ms, percentile(95%) = 16.9045 ms, percentile(99%) = 17.1071 ms\n",
    "[06/14/2023-08:26:15] [I] Enqueue Time: min = 0.503418 ms, max = 3.00732 ms, mean = 1.80985 ms, median = 1.79675 ms, percentile(90%) = 2.27783 ms, percentile(95%) = 2.5625 ms, percentile(99%) = 2.90527 ms\n",
    "[06/14/2023-08:26:15] [I] H2D Latency: min = 0.141357 ms, max = 0.233276 ms, mean = 0.17143 ms, median = 0.169434 ms, percentile(90%) = 0.1875 ms, percentile(95%) = 0.195068 ms, percentile(99%) = 0.206055 ms\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12b60890-4f72-4ed1-b8e1-9fa05ed44c7c",
   "metadata": {},
   "source": [
    "It's important to remember **that some methods may not show any improvements** for unknown reasons or unpredictable behavior. So, it's a good idea to carefully check and validate the results, comparing them with the original methods."
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}