pytorch_onnx_trt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define a model that serves as an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "class CNN(torch.nn.Module):\n",
    "    def __init__(self, num_classes=10,):\n",
    "        super(CNN, self).__init__()\n",
    "        self.layer1 = torch.nn.Conv2d(3,16,3)\n",
    "        self.layer2 = torch.nn.Conv2d(16,64,5)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        \n",
    "        # TAKE CARE HERE\n",
    "        # Ceil_mode must be False, because onnx eporter does NOT support ceil_mode=True\n",
    "        self.max_pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, ceil_mode=False) \n",
    "        self.avg_pool = torch.nn.AdaptiveAvgPool2d((1,1)) \n",
    "        \n",
    "        self.fc = torch.nn.Linear(64,num_classes)\n",
    "        self.batch_size_onnx = 0\n",
    "        # FLAG for output ONNX model\n",
    "        self.export_to_onnx_mode = False                  \n",
    "    def forward_onnx(self, X_in):\n",
    "        print(\"Function forward_onnx called! \\n\")\n",
    "        x = self.layer1(X_in)\n",
    "        x = self.relu(x)\n",
    "        x = self.max_pool(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.avg_pool(x)\n",
    "        assert self.batch_size_onnx > 0\n",
    "        length_of_fc_layer = 64 # For exporting an onnx model that fit the TensorRT, processes here should be DETERMINISITC!\n",
    "        x = x.view(self.batch_size_onnx, length_of_fc_layer) # \n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "       \n",
    "    def forward_default(self, X_in):\n",
    "        print(\"Function forward_default called! \\n\")\n",
    "        x = self.layer1(X_in)\n",
    "        x = self.relu(x)\n",
    "        x = self.max_pool(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.avg_pool(x)\n",
    "        \n",
    "        # Such an operationt is not deterministic since it would depend on the input and therefore would result in errors\n",
    "        length_of_fc_layer = x.size(1) \n",
    "        x = x.view(-1, length_of_fc_layer)\n",
    "        \n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "        \n",
    "    def __call__(self, *args,**kargs):\n",
    "        if self.export_to_onnx_mode:\n",
    "            return self.forward_onnx(*args,**kargs)\n",
    "        else:\n",
    "            return self.forward_default(*args,**kargs)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export a PyTorch model to an ONNX model\n",
    "TensorRT does not support to PyTorch model directly, but fortunately we could take advantage of ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch model saved to ./model.pth\n",
      "\n",
      "Function forward_onnx called! \n",
      "\n",
      "graph(%input : Float(2, 3, 64, 64),\n",
      "      %layer1.weight : Float(16, 3, 3, 3),\n",
      "      %layer1.bias : Float(16),\n",
      "      %layer2.weight : Float(64, 16, 5, 5),\n",
      "      %layer2.bias : Float(64),\n",
      "      %fc.weight : Float(10, 64),\n",
      "      %fc.bias : Float(10)):\n",
      "  %7 : Float(2, 16, 62, 62) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%input, %layer1.weight, %layer1.bias), scope: Conv2d\n",
      "  %8 : Float(2, 16, 62, 62) = onnx::Relu(%7), scope: ReLU\n",
      "  %9 : Float(2, 16, 60, 60) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%8), scope: MaxPool2d\n",
      "  %10 : Float(2, 64, 56, 56) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[5, 5], pads=[0, 0, 0, 0], strides=[1, 1]](%9, %layer2.weight, %layer2.bias), scope: Conv2d\n",
      "  %11 : Float(2, 64, 56, 56) = onnx::Relu(%10), scope: ReLU\n",
      "  %12 : Float(2, 64, 1, 1) = onnx::GlobalAveragePool(%11), scope: AdaptiveAvgPool2d\n",
      "  %13 : Float(2, 64) = onnx::Flatten[axis=1](%12)\n",
      "  %output : Float(2, 10) = onnx::Gemm[alpha=1, beta=1, transB=1](%13, %fc.weight, %fc.bias), scope: Linear\n",
      "  return (%output)\n",
      "\n",
      "ONNX model exported to ./model.onnx\n",
      "\n"
     ]
    }
   ],
   "source": [
    "onnx_model_path = './model.onnx'\n",
    "pth_model_path  = './model.pth'\n",
    "\n",
    "# Load the model\n",
    "model = CNN(10)\n",
    "# Save the state dict, which will be loaded later\n",
    "print('PyTorch model saved to {}\\n'.format(pth_model_path))\n",
    "torch.save(model.state_dict(), pth_model_path)\n",
    "\n",
    "# This is for ONNX exporter to track all the operations inside the model\n",
    "batch_size_of_dummy_input = 2 # Any size you want\n",
    "dummy_input = torch.zeros([batch_size_of_dummy_input, 3, 64, 64], dtype=torch.float32) \n",
    "\n",
    "model.batch_size_onnx = batch_size_of_dummy_input\n",
    "model.export_to_onnx_mode = True\n",
    "input_names = [ \"input\" ]\n",
    "output_names = [ \"output\"] # Multiple inputs and outputs are supported\n",
    "with torch.no_grad():\n",
    "    # If verbose is set to False. The information below won't displayed\n",
    "    torch.onnx.export(model, dummy_input, onnx_model_path, verbose=True, input_names=input_names, output_names=output_names)\n",
    "print('ONNX model exported to {}\\n'.format(onnx_model_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare some useful functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycuda.autoinit\n",
    "import numpy as np\n",
    "import pycuda.driver as cuda\n",
    "import tensorrt as trt\n",
    "import torch\n",
    "\n",
    "TRT_LOGGER = trt.Logger() # This logger is required to build an engine\n",
    "\n",
    "\n",
    "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__()\n",
    "    \n",
    "def allocate_buffers(engine):\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    bindings = []\n",
    "    stream = cuda.Stream()\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",
    "    return inputs, outputs, bindings, stream\n",
    "        \n",
    "def get_engine(max_batch_size=1, onnx_file_path=\"\", engine_file_path=\"\",\\\n",
    "               fp16_mode=False, int8_mode=False, save_engine=False,\n",
    "              ):\n",
    "    \"\"\"Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it.\"\"\"\n",
    "    def build_engine(max_batch_size, save_engine):\n",
    "        \"\"\"Takes an ONNX file and creates a TensorRT engine to run inference with\"\"\"\n",
    "        with trt.Builder(TRT_LOGGER) as builder, \\\n",
    "                builder.create_network() as network,\\\n",
    "                trt.OnnxParser(network, TRT_LOGGER) as parser:\n",
    "            \n",
    "            builder.max_workspace_size = 1 << 30 # Your workspace size\n",
    "            builder.max_batch_size = max_batch_size\n",
    "            #pdb.set_trace()\n",
    "            builder.fp16_mode = fp16_mode  # Default: False\n",
    "            builder.int8_mode = int8_mode  # Default: False\n",
    "            if int8_mode:\n",
    "                # To be updated\n",
    "                raise NotImplementedError\n",
    "                \n",
    "            # Parse model file\n",
    "            if not os.path.exists(onnx_file_path):\n",
    "                quit('ONNX file {} not found'.format(onnx_file_path))\n",
    "                \n",
    "            print('Loading ONNX file from path {}...'.format(onnx_file_path))\n",
    "            with open(onnx_file_path, 'rb') as model:\n",
    "                print('Beginning ONNX file parsing')\n",
    "                parser.parse(model.read())\n",
    "            print('Completed parsing of ONNX file')\n",
    "            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))\n",
    "            #pdb.set_trace()\n",
    "            #network.mark_output(network.get_layer(network.num_layers-1).get_output(0)) # Riz   \n",
    "            #network.mark_output(network.get_layer(network.num_layers-1).get_output(1)) # Riz\n",
    "            \n",
    "            engine = builder.build_cuda_engine(network)\n",
    "            if engine is None:\n",
    "                print(\"Failed to create the engine\")\n",
    "                return None\n",
    "            print(\"Completed creating Engine\")\n",
    "            \n",
    "            if save_engine:\n",
    "                with open(engine_file_path, \"wb\") as f:\n",
    "                    f.write(engine.serialize())\n",
    "            return engine\n",
    "        \n",
    "    if os.path.exists(engine_file_path):\n",
    "        # If a serialized engine exists, load it instead of building a new one.\n",
    "        print(\"Reading engine from file {}\".format(engine_file_path))\n",
    "        with open(engine_file_path, \"rb\") as f, trt.Runtime(TRT_LOGGER) as runtime:\n",
    "            return runtime.deserialize_cuda_engine(f.read())\n",
    "    else:\n",
    "        return build_engine(max_batch_size, save_engine)\n",
    "    \n",
    "def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):\n",
    "    # Transfer data from CPU 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, bindings=bindings, 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]\n",
    "\n",
    "def postprocess_the_outputs(h_outputs, shape_of_output):\n",
    "    h_outputs = h_outputs.reshape(*shape_of_output)\n",
    "    return h_outputs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load an ONNX model and build an TensorRT engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading ONNX file from path ./model.onnx...\n",
      "Beginning ONNX file parsing\n",
      "Completed parsing of ONNX file\n",
      "Building an engine from file ./model.onnx; this may take a while...\n",
      "Completed creating Engine\n",
      "Function forward_default called! \n",
      "\n",
      "Inference time with the TensorRT engine: 0.0009050369262695312\n",
      "Inference time with the PyTorch model: 0.0011165142059326172\n",
      "MSE Error = 5.2735593007950446e-17\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "onnx_model_path = './model.onnx'\n",
    "pytorch_model_path = './model.pth'\n",
    "\n",
    "# These two modes are dependent on hardwares\n",
    "fp16_mode = False\n",
    "int8_mode = False \n",
    "trt_engine_path = './model_fp16_{}_int8_{}.trt'.format(fp16_mode, int8_mode)\n",
    "\n",
    "\n",
    "max_batch_size = 1 # The batch size of input mush be smaller the max_batch_size once the engine is built\n",
    "x_input = np.random.rand(max_batch_size, 3, 64, 64).astype(dtype=np.float32) \n",
    "\n",
    "# Build an engine\n",
    "engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode, int8_mode)\n",
    "# Please check that engine is not None\n",
    "# If it is None, that means it failed to create the engine. \n",
    "# Create the context for this engine\n",
    "context = engine.create_execution_context() \n",
    "# Allocate buffers for input and output\n",
    "inputs, outputs, bindings, stream = allocate_buffers(engine) # input, output: host # bindings \n",
    "\n",
    "\n",
    "# Do inference\n",
    "shape_of_output = (max_batch_size, 10)\n",
    "# Load data to the buffer\n",
    "inputs[0].host = x_input.reshape(-1)\n",
    "# inputs[1].host = ... for multiple input\n",
    "t1 = time.time()\n",
    "trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) # numpy data\n",
    "t2 = time.time()\n",
    "output_from_trt_engine = postprocess_the_outputs(trt_outputs[0], shape_of_output)\n",
    "\n",
    "# Compare with the PyTorch\n",
    "pth_model = CNN(10)\n",
    "pth_model.load_state_dict(torch.load(pytorch_model_path))\n",
    "pth_model.cuda()\n",
    "\n",
    "x_input_pth = torch.from_numpy(x_input).cuda()\n",
    "pth_model.export_to_onnx_mode = False\n",
    "t3 = time.time()\n",
    "output_from_pytorch_model =  pth_model(x_input_pth)\n",
    "t4 = time.time()\n",
    "output_from_pytorch_model = output_from_pytorch_model.cpu().data.numpy()\n",
    "\n",
    "\n",
    "mse = np.mean((output_from_trt_engine - output_from_pytorch_model)**2)\n",
    "print(\"Inference time with the TensorRT engine: {}\".format(t2-t1))\n",
    "print(\"Inference time with the PyTorch model: {}\".format(t4-t3))\n",
    "print('MSE Error = {}'.format(mse))\n",
    "\n",
    "\n"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}