3_Convert_ONNX_to_TensorRT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3_Convert_ONNX_to_TensorRT\n",
    "\n",
    "This notebook will convert ONNX model to TensorRT serialized engine.\n",
    "\n",
    "The batch size is choosen to be 1.\n",
    "Conversion from ONNX to TensorRT would use the ONNXClassifierWrapper from Nvidia TensorRT repository.\n",
    "\n",
    "https://github.com/NVIDIA/TensorRT/tree/master/quickstart/IntroNotebooks\n",
    "\n",
    "Please see onnx_helper.py.\n",
    "Important modification to note is, explicit batch must be set.\n",
    "\n",
    "**explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)**\n",
    "\n",
    "We can also use the trtexec command to convert.\n",
    "\n",
    "trtexec --onnx=[ONNX file name.onnx] --saveEngine=[TensorRT engine.trt]  --explicitBatch\n",
    "\n",
    "### Jetson Nano Specific Information:\n",
    " In Jetson Nano, during the conversion, RAM might run out.\n",
    " Hence please set up swap space in the device.\n",
    " \n",
    " Please refer to the setup instructions [jetson_nano_setup_instructions.md] in this repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import packages\n",
    "import numpy as np\n",
    "from onnx_helper import ONNXClassifierWrapper,convert_onnx_to_engine\n",
    "import torch\n",
    "import json\n",
    "from PIL import Image\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set constants\n",
    "BATCH_SIZE=1\n",
    "N_CLASSES=1000\n",
    "PRECISION=np.float32\n",
    "image_size=224"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PyTorch transform\n",
    "tfms=transforms.Compose([transforms.Resize(image_size),\n",
    "                         transforms.CenterCrop(image_size),\n",
    "                         transforms.ToTensor(),\n",
    "                         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Img shape torch.Size([1, 3, 224, 224])\n",
      "(1, 3, 224, 224)\n"
     ]
    }
   ],
   "source": [
    "img=Image.open('panda.jpg')\n",
    "img=tfms(img)#.unsqueeze(0)\n",
    "\n",
    "img=torch.unsqueeze(img, 0)\n",
    "print(\"Img shape\",img.shape)\n",
    "BATCH_SIZE=1\n",
    "dummy_batch=np.zeros((BATCH_SIZE,3,224,224))\n",
    "for idx in range(BATCH_SIZE):\n",
    "    dummy_batch[idx]=img\n",
    "print(dummy_batch.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parsing ONNX file.\n",
      "Building TensorRT engine. This may take a few minutes.\n"
     ]
    }
   ],
   "source": [
    "ONNX_PATH='models/efficientnetb2_batch1.onnx'\n",
    "BATCH_SIZE=1\n",
    "TRT_PATH='models/efficientnetb2_batch1.trt'\n",
    "trt_engine=convert_onnx_to_engine(ONNX_PATH, TRT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}