Compare_trt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e3459299",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import time\n",
    "import onnx\n",
    "import torch\n",
    "import torchvision\n",
    "from PIL import Image\n",
    "\n",
    "# trt\n",
    "import numpy as np\n",
    "import tensorrt as trt\n",
    "import pycuda.autoinit\n",
    "import pycuda.driver as cuda\n",
    "\n",
    "from torch2trt import torch2trt\n",
    "from torch2trt import TRTModule"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d2416f6",
   "metadata": {},
   "source": [
    "### Export PyTorch model to ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0a84ea6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def torch_2_onnx(torch_model, savepath, batch_size=1, input_name=['input'], output_name=['output'], shape=[3, 224, 224]):\n",
    "    dummy_input = torch.randn(batch_size, *shape).cuda()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        torch.onnx.export(\n",
    "            torch_model,\n",
    "            dummy_input,\n",
    "            savepath,\n",
    "            input_names=input_name,\n",
    "            output_names=output_name,\n",
    "            verbose=False, # True for show INFO\n",
    "            opset_version=11\n",
    "        )\n",
    "        \n",
    "    print('ONNX model exported to {}\\n'.format(onnx_model_path))\n",
    "\n",
    "    # check\n",
    "    test = onnx.load(savepath)\n",
    "    onnx.checker.check_model(test)\n",
    "    print(\"==> Passed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "49e878d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_torch_model(model, pth_model_path):\n",
    "    print('PyTorch model saved to {}\\n'.format(pth_model_path))\n",
    "    torch.save(model.state_dict(), pth_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cf541c4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch model saved to model.pth\n",
      "\n",
      "ONNX model exported to model.onnx\n",
      "\n",
      "==> Passed\n"
     ]
    }
   ],
   "source": [
    "onnx_model_path = 'model.onnx'\n",
    "pth_model_path  = 'model.pth'\n",
    "\n",
    "# Load the model\n",
    "model = torchvision.models.resnet50(pretrained=True).eval().cuda()\n",
    "save_torch_model(model, pth_model_path)\n",
    "\n",
    "torch_2_onnx(model, onnx_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a722e61b",
   "metadata": {},
   "source": [
    "### Export ONNX model to TRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8cebb96e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trt version 7.2.3.4\n",
      "TRT model exported to model.trt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def onnx_2_trt(onnx_model_path, trt_model_path='model.trt'):\n",
    "    TRT_LOGGER = trt.Logger() # This logger is required to build an engine\n",
    "\n",
    "    EXPLICIT_BATCH = []\n",
    "    print('trt version', trt.__version__)\n",
    "    if trt.__version__[0] >= '7':\n",
    "        EXPLICIT_BATCH.append(\n",
    "            1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))\n",
    "\n",
    "    with trt.Builder(TRT_LOGGER) as builder,\\\n",
    "        builder.create_network(*EXPLICIT_BATCH) as network,\\\n",
    "        trt.OnnxParser(network, TRT_LOGGER) as parser:\n",
    "\n",
    "        builder.max_workspace_size = 1 << 28\n",
    "        builder.max_batch_size = 1\n",
    "\n",
    "        with open(onnx_model_path, 'rb') as f:\n",
    "            if not parser.parse(f.read()):\n",
    "                for error in range(parser.num_errors):\n",
    "                    print(parser.get_error(error))\n",
    "\n",
    "        # reshape input from 32 to 1\n",
    "        shape = list(network.get_input(0).shape)\n",
    "        engine = builder.build_cuda_engine(network)\n",
    "        with open(trt_model_path, 'wb') as f:\n",
    "            f.write(engine.serialize())\n",
    "    \n",
    "    print('TRT model exported to {}\\n'.format(trt_model_path))\n",
    "    \n",
    "onnx_2_trt(onnx_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef9671b5",
   "metadata": {},
   "source": [
    "### Inference with TRT Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "14792b7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HostDeviceMem(object):\n",
    "    def __init__(self, host_mem, device_mem):\n",
    "        \"\"\"Within this context, host_mom means the cpu memory and device means the GPU memory\n",
    "        \"\"\"\n",
    "        self.host = host_mem \n",
    "        self.device = device_mem\n",
    "    def __str__(self):\n",
    "        return \"Host:\\n\" + str(self.host) + \"\\nDevice:\\n\" + str(self.device)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return self.__str__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6ad2be91",
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):\n",
    "    \"\"\"do_inference (for TensorRT 6.x or lower)\n",
    "    This function is generalized for multiple inputs/outputs.\n",
    "    Inputs and outputs are expected to be lists of HostDeviceMem objects.\n",
    "    \"\"\"\n",
    "    # Transfer input data to the GPU.\n",
    "    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]\n",
    "    # Run inference.\n",
    "    context.execute_async(batch_size=batch_size,\n",
    "                          bindings=bindings,\n",
    "                          stream_handle=stream.handle)\n",
    "    # Transfer predictions back from the GPU.\n",
    "    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]\n",
    "    # Synchronize the stream\n",
    "    stream.synchronize()\n",
    "    # Return only the host outputs.\n",
    "    return [out.host for out in outputs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6a32c1df",
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_inference_v2(context, bindings, inputs, outputs, stream):\n",
    "    \"\"\"do_inference_v2 (for TensorRT 7.0+)\n",
    "    This function is generalized for multiple inputs/outputs for full\n",
    "    dimension networks.\n",
    "    Inputs and outputs are expected to be lists of HostDeviceMem objects.\n",
    "    \"\"\"\n",
    "    # Transfer input data to the GPU.\n",
    "    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]\n",
    "    # Run inference.\n",
    "    context.execute_async_v2(bindings=bindings,\n",
    "                             stream_handle=stream.handle)\n",
    "    # Transfer predictions back from the GPU.\n",
    "    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]\n",
    "    # Synchronize the stream\n",
    "    stream.synchronize()\n",
    "    # Return only the host outputs.\n",
    "    return [out.host for out in outputs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e179469c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def allocate_buffer(engine):\n",
    "    \"\"\"Allocates all host/device in/out buffers required for an engine.\"\"\"\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    bindings = []\n",
    "    output_idx = 0\n",
    "    stream = cuda.Stream()\n",
    "    \n",
    "    for binding in engine:\n",
    "        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size\n",
    "        dtype = trt.nptype(engine.get_binding_dtype(binding))\n",
    "        # Allocate host and device buffers\n",
    "        host_mem = cuda.pagelocked_empty(size, dtype)\n",
    "        device_mem = cuda.mem_alloc(host_mem.nbytes)\n",
    "        # Append the device buffer to device bindings.\n",
    "        bindings.append(int(device_mem))\n",
    "        # Append to the appropriate list.\n",
    "        if engine.binding_is_input(binding):\n",
    "            inputs.append(HostDeviceMem(host_mem, device_mem))\n",
    "        else:\n",
    "            outputs.append(HostDeviceMem(host_mem, device_mem))\n",
    "            \n",
    "    return inputs, outputs, bindings, stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "ae20ce55",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TRT_Engine(object):\n",
    "    \"\"\"\n",
    "    A TRT Engine Demo, include engine inintial and inference\n",
    "    \"\"\" \n",
    "    def __init__(self, trt_bin, cuda_ctx=None):\n",
    "        \"\"\"\n",
    "        param:\n",
    "            trt_bin: 'path/to/model.trt'\n",
    "            cuda_ctx: cuda.Device(0).make_context(), 0 is GPU number\n",
    "        \"\"\"\n",
    "        self.trt_bin = trt_bin\n",
    "        self.cuda_ctx = cuda_ctx\n",
    "        if self.cuda_ctx:\n",
    "            self.cuda_ctx.push()\n",
    "        \n",
    "        self.trt_logger = trt.Logger(trt.Logger.INFO)\n",
    "        self.engine = self._load_engine()\n",
    "        self.inference_fn = do_inference if trt.__version__[0] < '7' else do_inference_v2\n",
    "        \n",
    "        try:\n",
    "            self.context = self.engine.create_execution_context()\n",
    "            self.inputs, self.outputs, self.bindings, self.stream = allocate_buffer(self.engine)\n",
    "        except Exception as e:\n",
    "            raise RuntimeError('Fail to allocate CUDA resources') from e\n",
    "        finally:\n",
    "            if self.cuda_ctx:\n",
    "                self.cuda_ctx.pop()\n",
    "                \n",
    "    def _load_engine(self):\n",
    "        with open(self.trt_bin, 'rb')as f, trt.Runtime(self.trt_logger) as runtime:\n",
    "            return runtime.deserialize_cuda_engine(f.read())\n",
    "        \n",
    "    def __del__(self):\n",
    "        \"\"\"Free CUDA memories.\"\"\"\n",
    "        del self.outputs\n",
    "        del self.inputs\n",
    "        del self.stream\n",
    "        \n",
    "    def __call__(self, img):\n",
    "        self.inputs[0].host = np.ascontiguousarray(img)\n",
    "        \n",
    "        if self.cuda_ctx:\n",
    "            self.cuda_ctx.push()\n",
    "        \n",
    "        trt_outputs = self.inference_fn(\n",
    "            context = self.context,\n",
    "            bindings = self.bindings,\n",
    "            inputs = self.inputs,\n",
    "            outputs = self.outputs,\n",
    "            stream = self.stream\n",
    "        )\n",
    "        \n",
    "        if self.cuda_ctx:\n",
    "            self.cuda_ctx.pop()\n",
    "        \n",
    "        return trt_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "ffbd8498",
   "metadata": {},
   "outputs": [],
   "source": [
    "engine = TRT_Engine('model.trt', False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "d55ab17e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((3, 224, 224), (1, 3, 224, 224))"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img = cv2.imread('1.jpg')\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "img = cv2.resize(img, (224, 224))\n",
    "\n",
    "mean = np.array([0.485, 0.456, 0.406])\n",
    "std = np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "img_norm = (img / 255. - mean) / std\n",
    "img_norm = img_norm.transpose((2, 0, 1)).astype(np.float32)\n",
    "img_4dims = np.expand_dims(img_norm, axis=0)\n",
    "\n",
    "img_norm.shape, img_4dims.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "f49d3f64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.003329808712005615"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = 0\n",
    "for i in range(100):\n",
    "    t1 = time.time()\n",
    "    y_trt = engine(img_norm)\n",
    "    t2 = time.time()\n",
    "    s += t2 - t1\n",
    "s / 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "2ff53a88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 1000)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(y_trt).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b6a4bf6",
   "metadata": {},
   "source": [
    "### Export PyTorch to TRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "950126f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def torch_2_trt(model, x, trt_model_path):\n",
    "    assert isinstance(x, torch.cuda.FloatTensor), print('Invalid Input Type - {}'.format(type(x)))\n",
    "    model_trt = torch2trt(model, [x])\n",
    "    torch.save(model_trt.state_dict(), trt_model_path)\n",
    "    print('TRT model exported to {}\\n'.format(trt_model_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7beef529",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception ignored in: <function TRT_Engine.__del__ at 0x7f7fe63651f0>\n",
      "Traceback (most recent call last):\n",
      "  File \"/tmp/ipykernel_28441/798862869.py\", line 32, in __del__\n",
      "AttributeError: outputs\n",
      "Exception ignored in: <function TRT_Engine.__del__ at 0x7f7fe636bd30>\n",
      "Traceback (most recent call last):\n",
      "  File \"/tmp/ipykernel_28441/3935465325.py\", line 32, in __del__\n",
      "AttributeError: outputs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 224, 224])\n",
      "TRT model exported to model_trt.pth\n",
      "\n"
     ]
    }
   ],
   "source": [
    "x = torch.from_numpy(img_4dims).float().cuda()\n",
    "print(x.size())\n",
    "torch_2_trt(model, x, 'model_trt.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2baa2504",
   "metadata": {},
   "source": [
    "### Inference with torch_2_TRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "740ff432",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trt1 = torch2trt(model, [x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "b473b345",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.003029191493988037"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = 0\n",
    "for i in range(100):\n",
    "    t1 = time.time()\n",
    "    y1 = model_trt1(x)\n",
    "    t2 = time.time()\n",
    "    s += t2 - t1\n",
    "s / 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "09493835",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_trt2 = TRTModule()\n",
    "model_trt2.load_state_dict(torch.load('model_trt.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "78896cfc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0027715039253234865"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = 0\n",
    "for i in range(100):\n",
    "    t3 = time.time()\n",
    "    y2 = model_trt2(x)\n",
    "    t4 = time.time()\n",
    "    s += t4 - t3\n",
    "s / 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "2711264a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.006194736957550049"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = 0\n",
    "for i in range(100):\n",
    "    t3 = time.time()\n",
    "    y3 = model(x)\n",
    "    t4 = time.time()\n",
    "    s += t4 - t3\n",
    "s / 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "7ddcbb47",
   "metadata": {},
   "outputs": [],
   "source": [
    "y1 = y1.detach().cpu()\n",
    "y2 = y2.detach().cpu()\n",
    "y3 = y3.detach().cpu()\n",
    "y_trt = torch.from_numpy(np.array(y_trt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "7fb43559",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(1.0541e-11), tensor(1.0549e-11), tensor(1.3717e-13))"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.mean(torch.pow(y3 - y2, 2)), torch.mean(torch.pow(y3 - y1, 2)), torch.mean(torch.pow(y2 - y1, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "d844acde",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(1.0541e-11), tensor(0.), tensor(1.3717e-13))"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.mean(torch.pow(y3 - y_trt, 2)), torch.mean(torch.pow(y2 - y_trt, 2)), torch.mean(torch.pow(y1 - y_trt, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f738868c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "290e07b8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a27d13",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d57ec3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1b55894",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee7ae8a6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25a9dca1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fad13ba",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de13b56e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4bf59c2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}