03_Infer_using_TensorRT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "from skimage.io import imread, imshow\n",
    "from skimage.transform import resize\n",
    "import tensorrt as trt; print('TensorRT Version: {}'.format(trt.__version__))\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set location of base directory\n",
    "base_path = os.path.join('/', 'workspace', 'optimization')\n",
    "\n",
    "# Set location of datasets\n",
    "datasets_path = os.path.join(base_path, 'datasets')\n",
    "images_path = os.path.join(datasets_path, 'images')\n",
    "masks_path = os.path.join(datasets_path, 'masks')\n",
    "\n",
    "# Set location of model\n",
    "artifacts_path = os.path.join(base_path, 'artifacts')\n",
    "engine_file_name = 'Jan_2019_99_w_rejects_{}_b{}_{}.engine'\n",
    "engine_file_path = os.path.join(artifacts_path, engine_file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set network settings\n",
    "n_channel, n_height, n_width = 1, 512, 384\n",
    "dimensions = [n_channel, n_height, n_width]\n",
    "batch_size = 1\n",
    "precision = 'fp16'  # options are 'fp16' (default), 'int8', and 'fp32'\n",
    "architecture = 'v100'  # options are 't4' (default), 'v100' and 'xavier'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set the logger severity higher to suppress messages (or lower to display more messages).\n",
    "# TRT_LOGGER = trt.Logger(trt.Logger.WARNING)\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create runtime\n",
    "runtime = trt.Runtime(TRT_LOGGER)\n",
    "print(runtime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read engine\n",
    "engine_file_path = engine_file_path.format(architecture, batch_size, precision)\n",
    "with open(engine_file_path, 'rb') as f:\n",
    "    engine = runtime.deserialize_cuda_engine(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check dimensions\n",
    "input_dimensions = engine.get_binding_shape(0)\n",
    "output_dimensions = engine.get_binding_shape(1)\n",
    "print('Input Dimensions:', input_dimensions)\n",
    "print('Output Dimensions:', output_dimensions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create execution context\n",
    "context = engine.create_execution_context()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load test images and masks\n",
    "images_file_names = os.listdir(images_path)\n",
    "images_file_paths = [os.path.join(images_path, f) for f in images_file_names]\n",
    "masks_file_names = os.listdir(masks_path)\n",
    "masks_file_paths = [os.path.join(masks_path, f) for f in masks_file_names]\n",
    "\n",
    "# Print first 2\n",
    "print('Images:', images_file_paths[:2])\n",
    "print('Masks:', masks_file_paths[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the jth image and mask\n",
    "j = 0\n",
    "image_file_path = images_file_paths[j]\n",
    "mask_file_path = masks_file_paths[j]\n",
    "print('Image:', image_file_path)\n",
    "print('Mask:', mask_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load image and mask\n",
    "image = imread(image_file_path)\n",
    "mask = imread(mask_file_path)\n",
    "print('Image:', image.shape, image.dtype)\n",
    "print('Mask:', mask.shape, mask.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change image dimension from 3 to 1\n",
    "image = image[:, :, 0]\n",
    "print('Image:', image.shape, image.dtype)\n",
    "print('Mask:', mask.shape, mask.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show image\n",
    "imshow(image, cmap='binary'); plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show mask\n",
    "imshow(mask, cmap='binary'); plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resize image\n",
    "image = resize(image, (n_height, n_width))\n",
    "print('Image:', image.shape, image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add channel dimension\n",
    "image = image[np.newaxis, :, :]\n",
    "print('Image:', image.shape, image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change type\n",
    "image = image.astype(np.float32)\n",
    "print('Image:', image.shape, image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add batch dimension\n",
    "image = image[np.newaxis, :, :, :]\n",
    "print('Image:', image.shape, image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Repeat data along batch axis\n",
    "batch_size = 1\n",
    "image = np.repeat(image, batch_size, axis=0)\n",
    "print('Image:', image.shape, image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create new data\n",
    "# new_data = image / 255.\n",
    "new_data = image\n",
    "output_dimensions = (batch_size, n_height, n_width, n_channel)\n",
    "output = np.empty(output_dimensions, dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allocate device memory\n",
    "d_input = cuda.mem_alloc(1 * new_data.nbytes)\n",
    "d_output = cuda.mem_alloc(1 * output.nbytes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create bindings and stream\n",
    "bindings = [int(d_input), int(d_output)]\n",
    "stream = cuda.Stream()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transfer input data to device\n",
    "cuda.memcpy_htod_async(d_input, new_data, stream)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute model\n",
    "context.execute_async(1, bindings, stream.handle, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transfer predictions back\n",
    "cuda.memcpy_dtoh_async(output, d_output, stream)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Syncronize threads\n",
    "stream.synchronize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Log info\n",
    "print(\"Prediction Shape: {}\".format(output.shape)) \n",
    "# print(\"Prediction: {} \".format(output))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show prediction\n",
    "if batch_size > 1:\n",
    "    imshow(np.squeeze(output[0]), cmap='binary')\n",
    "else:\n",
    "    imshow(np.squeeze(output), cmap='binary')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 2000\n",
    "start = time.time()\n",
    "for _ in range(n):\n",
    "    # Transfer input data to device\n",
    "    cuda.memcpy_htod_async(d_input, new_data, stream)\n",
    "    \n",
    "    # Execute model\n",
    "    context.execute_async(1, bindings, stream.handle, None)\n",
    "    \n",
    "    # Transfer predictions back\n",
    "    cuda.memcpy_dtoh_async(output, d_output, stream)\n",
    "    \n",
    "    # Syncronize threads\n",
    "    stream.synchronize()\n",
    "end = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = end - start\n",
    "average_latency = delta / n\n",
    "average_throughput = batch_size * (1 / average_latency)\n",
    "print('Inference: {} seconds'.format(delta))\n",
    "print('Number of Inferences: {}'.format(n))\n",
    "print('Average Latency: {} seconds'.format(average_latency))\n",
    "print('Average Throughput w/ Batch Size {}: {} examples per second'.format(batch_size, average_throughput))"
   ]
  },
  {
   "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
}