02_Convert_UFF_to_TensorRT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Thu Feb  7 00:09:25 2019       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 418.25       Driver Version: 418.25       CUDA Version: 10.1     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM3...  On   | 00000000:34:00.0 Off |                    0 |\n",
      "| N/A   33C    P0    53W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   1  Tesla V100-SXM3...  On   | 00000000:36:00.0 Off |                    0 |\n",
      "| N/A   32C    P0    52W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   2  Tesla V100-SXM3...  On   | 00000000:39:00.0 Off |                    0 |\n",
      "| N/A   36C    P0    53W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   3  Tesla V100-SXM3...  On   | 00000000:3B:00.0 Off |                    0 |\n",
      "| N/A   35C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   4  Tesla V100-SXM3...  On   | 00000000:57:00.0 Off |                    0 |\n",
      "| N/A   33C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   5  Tesla V100-SXM3...  On   | 00000000:59:00.0 Off |                    0 |\n",
      "| N/A   35C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   6  Tesla V100-SXM3...  On   | 00000000:5C:00.0 Off |                    0 |\n",
      "| N/A   32C    P0    52W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   7  Tesla V100-SXM3...  On   | 00000000:5E:00.0 Off |                    0 |\n",
      "| N/A   36C    P0    52W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   8  Tesla V100-SXM3...  On   | 00000000:B7:00.0 Off |                    0 |\n",
      "| N/A   31C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   9  Tesla V100-SXM3...  On   | 00000000:B9:00.0 Off |                    0 |\n",
      "| N/A   31C    P0    49W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|  10  Tesla V100-SXM3...  On   | 00000000:BC:00.0 Off |                    0 |\n",
      "| N/A   33C    P0    52W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|  11  Tesla V100-SXM3...  On   | 00000000:BE:00.0 Off |                    0 |\n",
      "| N/A   33C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|  12  Tesla V100-SXM3...  On   | 00000000:E0:00.0 Off |                    0 |\n",
      "| N/A   32C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|  13  Tesla V100-SXM3...  On   | 00000000:E2:00.0 Off |                    0 |\n",
      "| N/A   31C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|  14  Tesla V100-SXM3...  On   | 00000000:E5:00.0 Off |                    0 |\n",
      "| N/A   35C    P0    53W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|  15  Tesla V100-SXM3...  On   | 00000000:E7:00.0 Off |                    0 |\n",
      "| N/A   34C    P0    51W / 350W |      0MiB / 32480MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                       GPU Memory |\n",
      "|  GPU       PID   Type   Process name                             Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow Version: 1.12.0\n",
      "TensorRT Version: 5.0.2.6\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "import tensorflow as tf; print('TensorFlow Version: {}'.format(tf.__version__))\n",
    "import tensorrt as trt; print('TensorRT Version: {}'.format(trt.__version__))\n",
    "import uff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = os.path.join('/', 'datasets', 'retail', 'iva', 'industrial-inspection', 'pg-defect-segmentation')\n",
    "models_dir = os.path.join(base_dir, 'models')\n",
    "artifacts_dir = os.path.join(models_dir, 'artifacts')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# network settings\n",
    "model_path = os.path.join(artifacts_dir, 'frozen_graph.uff')\n",
    "input_names = ['input_1']\n",
    "output_names = ['conv2d_19_2/Sigmoid']\n",
    "# n_channel, n_height, n_width = 1, 512, 512\n",
    "n_channel, n_height, n_width = 3, 512, 512\n",
    "# n_channel, n_height, n_width = 1, 512, 1136\n",
    "# n_channel, n_height, n_width = 1, 1024, 1024\n",
    "# n_channel, n_height, n_width = 1, 1024, 2272\n",
    "# n_channel, n_height, n_width = 1, 2048, 2048\n",
    "# n_channel, n_height, n_width = 1, 4096, 4096\n",
    "dimensions = [n_channel, n_height, n_width]\n",
    "batch_size = 1\n",
    "precision = 'fp16'  # options are 'fp16' (default), 'int8', and 'fp32'\n",
    "architecture = 't4'  # options are 't4' (default) and 'xavier'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_engine(filename):\n",
    "    # For more information on TRT basics, refer to the introductory samples.\n",
    "    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:\n",
    "        # Set max workspace\n",
    "        builder.max_workspace_size = 1 << 30\n",
    "        # Parse the Uff Network\n",
    "        parser.register_input(input_names[0], dimensions)\n",
    "        parser.register_output(output_names[0])\n",
    "        parser.parse(filename, network)\n",
    "        # Set precision\n",
    "        if precision == 'fp16':\n",
    "            builder.fp16_mode = True\n",
    "        elif precision == 'int8':\n",
    "            builder.int8_mode = True\n",
    "        # Set batch size\n",
    "        #builder.max_batch_size = batch_size\n",
    "        # Build and return an engine.\n",
    "        return builder.build_cuda_engine(network)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<tensorrt.tensorrt.ICudaEngine object at 0x7fe2d9ce88b8>\n"
     ]
    }
   ],
   "source": [
    "engine = build_engine(model_path)\n",
    "print(engine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing: /datasets/retail/iva/industrial-inspection/pg-defect-segmentation/models/artifacts/defect_classifier_t4_b1_fp16.engine\n"
     ]
    }
   ],
   "source": [
    "# write engine\n",
    "engine_path = os.path.join(artifacts_dir, 'defect_classifier_{}_b{}_{}.engine')\n",
    "engine_path = engine_path.format(architecture, batch_size, precision)\n",
    "with open(engine_path, 'wb') as file:\n",
    "    print('Writing: {}'.format(engine_path))\n",
    "    file.write(engine.serialize())"
   ]
  },
  {
   "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
}