{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechRecognition/Jasper/notebooks/Colab_Jasper_TRT_inference_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Gwt7z7qdmTbW"
   },
   "outputs": [],
   "source": [
    "# Copyright 2019 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "i4NKCp2VmTbn"
   },
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Jasper Inference Demo with NVIDIA TensorRT on Google Colab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fW0OKDzvmTbt"
   },
   "source": [
    "## Overview\n",
    "\n",
    "\n",
    "In this notebook, we will demo the process of carrying out inference on new audio segment using a pre-trained Pytorch Jasper model downloaded from the NVIDIA NGC Model registry with TensorRT (TRT). NVIDIA TensorRT is a platform for high-performance deep learning inference. It includes a deep learning inference optimizer and runtime that delivers low latency and high-throughput for deep learning inference applications. After optimizing the compute-intensive acoustic model with NVIDIA TensorRT, inference throughput increased by up to 1.8x over native PyTorch.\n",
    "\n",
    "The Jasper model is an end-to-end neural acoustic model for automatic speech recognition (ASR) that provides near state-of-the-art results on LibriSpeech among end-to-end ASR models without any external data. The Jasper architecture of convolutional layers was designed to facilitate fast GPU inference, by allowing whole sub-blocks to be fused into a single GPU kernel. This is important for meeting strict real-time requirements of ASR systems in deployment.The results of the acoustic model are combined with the results of external language models to get the top-ranked word sequences corresponding to a given audio segment. This post-processing step is called decoding.\n",
    "\n",
    "The original paper is Jasper: An End-to-End Convolutional Neural Acoustic Model https://arxiv.org/pdf/1904.03288.pdf.\n",
    "\n",
    "### Model architecture\n",
    "By default the model configuration is Jasper 10x5 with dense residuals. A Jasper BxR model has B blocks, each consisting of R repeating sub-blocks.\n",
    "Each sub-block applies the following operations in sequence: 1D-Convolution, Batch Normalization, ReLU activation, and Dropout. \n",
    "In the original paper Jasper is trained with masked convolutions, which masks out the padded part of an input sequence in a batch before the 1D-Convolution.\n",
    "For inference masking is not used. The reason for this is that in inference, the original mask operation does not achieve better accuracy than without the mask operation on the test and development dataset. However, no masking achieves better inference performance especially after TensorRT optimization.\n",
    "More information on the model architecture can be found in the [root folder](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/Jasper)\n",
    "\n",
    "### TensorRT Inference pipeline\n",
    "The Jasper inference pipeline consists of 3 components: data preprocessor, acoustic model and greedy decoder. The acoustic model is the most compute intensive, taking more than 90% of the entire end-to-end pipeline. The acoustic model is the only component with learnable parameters and also what differentiates Jasper from the competition. So, we focus on the acoustic model for the most part.\n",
    "For the non-TRT Jasper inference pipeline, all 3 components are implemented and run with native PyTorch. For the TensorRT inference pipeline, we show the speedup of running the acoustic model with TensorRT, while preprocessing and decoding are reused from the native PyTorch pipeline.\n",
    "To run a model with TensorRT, we first construct the model in PyTorch, which is then exported into an ONNX file. Finally, a TensorRT engine is constructed from the ONNX file, serialized to TRT plan file, and also launched to do inference.\n",
    "Note that TensorRT engine is being runtime optimized before serialization. TRT tries a vast set of options to find the strategy that performs best on user’s GPU - so it takes a few minutes. After the TRT plan file is created, it can be reused.\n",
    "\n",
    "\n",
    "### Requirement\n",
    "1. Before running this notebook, please set the Colab runtime environment to GPU via the menu *Runtime => Change runtime type => GPU*.\n",
    "\n",
    "For TRT FP16 and INT8 inference, an NVIDIA Volta, Turing or newer GPU generations is required. On Google Colab, this normally means a T4 GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 316
    },
    "colab_type": "code",
    "id": "HVsrGkj4Zn2L",
    "outputId": "7061e16f-7d7c-4be9-cba4-4c8f67fdafb0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wed Oct  2 02:42:12 2019       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 430.40       Driver Version: 418.67       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 K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
      "| N/A   60C    P0    70W / 149W |     69MiB / 11441MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                       GPU Memory |\n",
      "|  GPU       PID   Type   Process name                             Usage      |\n",
      "|=============================================================================|\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pV3rzgO8-tSK"
   },
   "source": [
    "The below code check whether a Tensor core GPU is present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "colab_type": "code",
    "id": "Djyvo8mm9poq",
    "outputId": "4ec6cda8-1e68-40b9-de79-54e4285f0b7a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor Core GPU Present: None\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "\n",
    "def check_tensor_core_gpu_present():\n",
    "    local_device_protos = device_lib.list_local_devices()\n",
    "    for line in local_device_protos:\n",
    "        if \"compute capability\" in str(line):\n",
    "            compute_capability = float(line.physical_device_desc.split(\"compute capability: \")[-1])\n",
    "            if compute_capability>=7.0:\n",
    "                return True\n",
    "    \n",
    "print(\"Tensor Core GPU Present:\", check_tensor_core_gpu_present())\n",
    "tensor_core_gpu = check_tensor_core_gpu_present()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FCEfkBAbbaLI"
   },
   "source": [
    "2. Next, we clone the NVIDIA Github Deep Learning Example repository and set up the workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 148
    },
    "colab_type": "code",
    "id": "y3u_VMjXtAto",
    "outputId": "84f4b72c-71cd-4415-d94f-1b87266eab74"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'DeepLearningExamples'...\n",
      "remote: Enumerating objects: 110, done.\u001b[K\n",
      "remote: Counting objects: 100% (110/110), done.\u001b[K\n",
      "remote: Compressing objects: 100% (90/90), done.\u001b[K\n",
      "remote: Total 4049 (delta 65), reused 35 (delta 17), pack-reused 3939\u001b[K\n",
      "Receiving objects: 100% (4049/4049), 32.29 MiB | 26.48 MiB/s, done.\n",
      "Resolving deltas: 100% (1875/1875), done.\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/NVIDIA/DeepLearningExamples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "colab_type": "code",
    "id": "-rE46y-ftAuQ",
    "outputId": "1bd44631-d98f-40a2-a999-3cc861c63d92"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/content/DeepLearningExamples/PyTorch/SpeechRecognition/Jasper/notebooks\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "WORKSPACE_DIR='/content/DeepLearningExamples/PyTorch/SpeechRecognition/Jasper/notebooks'\n",
    "os.chdir(WORKSPACE_DIR)\n",
    "print (os.getcwd())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yBeZjO4JtAwL"
   },
   "source": [
    "## Install NVIDIA TensorRT\n",
    "\n",
    "We will need to install NVIDIA TensorRT 6.0 runtime environment on Colab. First, check the Colab CUDA installed version. As of 2nd Oct 2019, `cuda-10.0` is the CUDA version on Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "id": "LfygzbP1Lz2b",
    "outputId": "a0c985a4-c6f8-47f7-a37a-a5c7e5168fef"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bin   cuda-10.0  games\t  lib\t       man   setup.cfg\tsrc\n",
      "cuda  etc\t include  LICENSE.txt  sbin  share\txgboost\n"
     ]
    }
   ],
   "source": [
    "!ls /usr/local/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yl4PyTsnQP3U"
   },
   "source": [
    "Next, we will need to install the NVIDIA TensorRT version that match the current Colab CUDA version, following the instruction at https://docs.nvidia.com/deeplearning/sdk/tensorrt-install-guide/index.html#maclearn-net-repo-install."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 935
    },
    "colab_type": "code",
    "id": "3WA9N43UTq_c",
    "outputId": "13db6406-2eea-45a0-d3b2-222b7668f951"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb\n",
    "\n",
    "dpkg -i nvidia-machine-learning-repo-*.deb\n",
    "apt-get update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XPsJhzqLXDAm"
   },
   "source": [
    "When using the NVIDIA Machine Learning network repository, Ubuntu will be default install TensorRT for the latest CUDA version. The following commands will install libnvinfer6 for an older CUDA version and hold the libnvinfer6 package at this version. Replace 6.0.1 with your version of TensorRT and cuda10.0 with your CUDA version for your Colab environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "EbF-JWGfK9Lo",
    "outputId": "41268923-4997-4f92-80bb-c718b6e56a64"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "version=\"6.0.1-1+cuda10.0\"\n",
    "sudo apt-get install libnvinfer6=${version} libnvonnxparsers6=${version} libnvparsers6=${version} libnvinfer-plugin6=${version} libnvinfer-dev=${version} libnvonnxparsers-dev=${version} libnvparsers-dev=${version} libnvinfer-plugin-dev=${version} python-libnvinfer=${version} python3-libnvinfer=${version}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 243
    },
    "colab_type": "code",
    "id": "zZO9cPHHLGoc",
    "outputId": "3f4e1307-5fae-4052-993c-0f237b6d4957"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "libnvinfer6 set on hold.\n",
      "libnvonnxparsers6 set on hold.\n",
      "libnvparsers6 set on hold.\n",
      "libnvinfer-plugin6 set on hold.\n",
      "libnvinfer-dev set on hold.\n",
      "libnvonnxparsers-dev set on hold.\n",
      "libnvparsers-dev set on hold.\n",
      "libnvinfer-plugin-dev set on hold.\n",
      "python-libnvinfer set on hold.\n",
      "python3-libnvinfer set on hold.\n",
      "W: Target Packages (Packages) is configured multiple times in /etc/apt/sources.list.d/nvidia-machine-learning.list:1 and /etc/apt/sources.list.d/nvidia-ml.list:1\n"
     ]
    }
   ],
   "source": [
    "!sudo apt-mark hold libnvinfer6 libnvonnxparsers6 libnvparsers6 libnvinfer-plugin6 libnvinfer-dev libnvonnxparsers-dev libnvparsers-dev libnvinfer-plugin-dev python-libnvinfer python3-libnvinfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 224
    },
    "colab_type": "code",
    "id": "wOo7YbuhLcUU",
    "outputId": "8497ea38-9a86-41a2-d68c-a2310c2a1899"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi  libnvinfer-dev                          6.0.1-1+cuda10.0                                  amd64        TensorRT development libraries and headers\n",
      "hi  libnvinfer-plugin-dev                   6.0.1-1+cuda10.0                                  amd64        TensorRT plugin libraries\n",
      "hi  libnvinfer-plugin6                      6.0.1-1+cuda10.0                                  amd64        TensorRT plugin libraries\n",
      "hi  libnvinfer6                             6.0.1-1+cuda10.0                                  amd64        TensorRT runtime libraries\n",
      "hi  libnvonnxparsers-dev                    6.0.1-1+cuda10.0                                  amd64        TensorRT ONNX libraries\n",
      "hi  libnvonnxparsers6                       6.0.1-1+cuda10.0                                  amd64        TensorRT ONNX libraries\n",
      "hi  libnvparsers-dev                        6.0.1-1+cuda10.0                                  amd64        TensorRT parsers libraries\n",
      "hi  libnvparsers6                           6.0.1-1+cuda10.0                                  amd64        TensorRT parsers libraries\n",
      "hi  python-libnvinfer                       6.0.1-1+cuda10.0                                  amd64        Python bindings for TensorRT\n",
      "hi  python3-libnvinfer                      6.0.1-1+cuda10.0                                  amd64        Python 3 bindings for TensorRT\n"
     ]
    }
   ],
   "source": [
    "!dpkg -l | grep TensorRT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XRMZiFjdUPCZ"
   },
   "source": [
    "A successful TensorRT installation should look like:\n",
    "\n",
    "```\n",
    "hi  libnvinfer-dev                          6.0.1-1+cuda10.0                                  amd64        TensorRT development libraries and headers\n",
    "hi  libnvinfer-plugin-dev                   6.0.1-1+cuda10.0                                  amd64        TensorRT plugin libraries\n",
    "hi  libnvinfer-plugin6                      6.0.1-1+cuda10.0                                  amd64        TensorRT plugin libraries\n",
    "hi  libnvinfer6                             6.0.1-1+cuda10.0                                  amd64        TensorRT runtime libraries\n",
    "hi  libnvonnxparsers-dev                    6.0.1-1+cuda10.0                                  amd64        TensorRT ONNX libraries\n",
    "hi  libnvonnxparsers6                       6.0.1-1+cuda10.0                                  amd64        TensorRT ONNX libraries\n",
    "hi  libnvparsers-dev                        6.0.1-1+cuda10.0                                  amd64        TensorRT parsers libraries\n",
    "hi  libnvparsers6                           6.0.1-1+cuda10.0                                  amd64        TensorRT parsers libraries\n",
    "hi  python-libnvinfer                       6.0.1-1+cuda10.0                                  amd64        Python bindings for TensorRT\n",
    "hi  python3-libnvinfer                      6.0.1-1+cuda10.0                                  amd64        Python 3 bindings for TensorRT\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lwpwNNw8FS8F"
   },
   "source": [
    "## Download pretrained Jasper model from NVIDIA GPU Cloud model repository\n",
    "\n",
    "NVIDIA provides pretrained Jasper models along with many other deep learning models such as ResNet, BERT, Transformer, SSD... at https://ngc.nvidia.com/catalog/models. Here, we will download and unzip pretrained Jasper Pytorch models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "np8WaN_FFaTF",
    "outputId": "34801d4c-a621-4e13-9f69-16b429ebeccf"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%%bash \n",
    "wget -nc -q --show-progress -O jasper_model.zip \\\n",
    "https://api.ngc.nvidia.com/v2/models/nvidia/jasperpyt_fp16/versions/1/zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "id": "gsJBwUgXHEkE",
    "outputId": "8e2214fc-6068-498e-aa73-9564642768ec"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Archive:  ./jasper_model.zip\n",
      "  inflating: jasper_fp16.pt          \n"
     ]
    }
   ],
   "source": [
    "!unzip -o ./jasper_model.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hbafoHBMXr0E"
   },
   "source": [
    "After a successful download, a Pytorch checkpoint named ` jasper_fp16.pt` should exist in the current notebooks directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "colab_type": "code",
    "id": "YC2Fu9rWG70U",
    "outputId": "8d8d8c9d-2ac3-454e-e2ae-72dcd8d56b81"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rw-r--r-- 1 root root 2661855989 Sep 10 00:33 jasper_fp16.pt\n"
     ]
    }
   ],
   "source": [
    "!ls -l jasper_fp16.pt  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "z-1-63MOG0um"
   },
   "source": [
    "## Install extra dependencies\n",
    "\n",
    "Before proceeding to creating the TensorRT execution engine from the Pytorch checkpoint, we shall install some extra dependency to load and convert the Pytorch model and process input audio files.\n",
    "\n",
    "- [Apex](https://nvidia.github.io/apex/): this is NVIDIA libraries for automatic mixed precision training in Pytorch\n",
    "- [Onnx](https://github.com/onnx/onnx): for processing ONNX model.\n",
    "- unidecode, soundfile, toml, pycuda: miscellaneous helper libraries\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "A6QC4ngRHw3a",
    "outputId": "ba7195a5-b235-47aa-8565-b671e40af831"
   },
   "outputs": [],
   "source": [
    "%%bash \n",
    "pip uninstall -y apex\n",
    "git clone https://www.github.com/nvidia/apex\n",
    "cd apex\n",
    "python setup.py install\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 729
    },
    "colab_type": "code",
    "id": "QPJbKeigIdOC",
    "outputId": "97f9cb47-e8ea-407e-e252-0ff46684a9e2"
   },
   "outputs": [],
   "source": [
    "!pip install unidecode soundfile toml pycuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 243
    },
    "colab_type": "code",
    "id": "UfMZMGbEMFXE",
    "outputId": "0724ae09-b195-4354-d16a-76d01ea81f82"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting onnx\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f5/f4/e126b60d109ad1e80020071484b935980b7cce1e4796073aab086a2d6902/onnx-1.6.0-cp36-cp36m-manylinux1_x86_64.whl (4.8MB)\n",
      "\u001b[K     |████████████████████████████████| 4.8MB 27kB/s \n",
      "\u001b[?25hRequirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from onnx) (1.12.0)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from onnx) (1.16.5)\n",
      "Collecting typing-extensions>=3.6.2.1 (from onnx)\n",
      "  Downloading https://files.pythonhosted.org/packages/27/aa/bd1442cfb0224da1b671ab334d3b0a4302e4161ea916e28904ff9618d471/typing_extensions-3.7.4-py3-none-any.whl\n",
      "Requirement already satisfied: protobuf in /usr/local/lib/python3.6/dist-packages (from onnx) (3.7.1)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf->onnx) (41.2.0)\n",
      "Installing collected packages: typing-extensions, onnx\n",
      "Successfully installed onnx-1.6.0 typing-extensions-3.7.4\n"
     ]
    }
   ],
   "source": [
    "!pip install onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bup2MbIjStm2"
   },
   "source": [
    "## Play with audio examples\n",
    "\n",
    "You can perform inference using pre-trained checkpoints which takes audio file (in .wav format) as input, and produces the corresponding text file. You can customize the content of the input .wav file. For example, there are several examples of input files at \"notebooks\" dirctory and we can listen to example1.wav:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 61
    },
    "colab_type": "code",
    "id": "u7J2WjikSxgu",
    "outputId": "3a862d2a-7fa6-4a65-9e84-817af265d74e"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/x-wav;base64,[encoded data, 429288 characters]\" type=\"audio/x-wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 19,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import IPython.display as ipd\n",
    "ipd.Audio('./example1.wav', rate=22050)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "la9uEY8ja9iV"
   },
   "source": [
    "You can also download your own audio sample to Colab with\n",
    "\n",
    "```!wget <link-to-.wav-file>```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "571hMhhGRnOB"
   },
   "source": [
    "## FP32 Inference with TensorRT\n",
    "\n",
    "\n",
    "### Creating TensorRT FP32 execution plan\n",
    "\n",
    "You can run inference using the trt/perf.py script:\n",
    "* the checkpoint is passed as `--ckpt` argument \n",
    "* `--model_toml` specifies the path to network configuration file (see examples in \"config\" directory)\n",
    "* `--make_onnx` exports to ONNX file at the path if set\n",
    "* `--engine_path` saves the engine file (*.plan) \n",
    "\n",
    "To create a new engine file (jasper.plan) for TensorRT and run it using fp32 (building the engine for the first time can take several minutes):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 224
    },
    "colab_type": "code",
    "id": "aJrN9pmdG4C8",
    "outputId": "35ce0bd5-ead4-4ab1-fd50-8f7062c14e77"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 1331142656 bytes == 0x15c680000 @  0x7f5e9070c887 0x7f5e8f002bf9 0x7f5e8f003acb 0x7f5e8f003b84 0x7f5e8f003f6c 0x7f5e4a95216f 0x7f5e4a9523f4 0x7f5e4053a411 0x7f5e8684837d 0x7f5e8657def4 0x56204c 0x4f88ba 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f7a28 0x4f876d\n",
      "tcmalloc: large alloc 1331142656 bytes == 0x1abbfa000 @  0x7f5e9070a1e7 0x5a1c5c 0x7f5e868486da 0x7f5e8657def4 0x56204c 0x4f88ba 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023\n",
      "[libprotobuf WARNING google/protobuf/io/coded_stream.cc:604] Reading dangerously large protocol message.  If the message turns out to be larger than 2147483647 bytes, parsing will be halted for security reasons.  To increase the limit (or to disable these warnings), see CodedInputStream::SetTotalBytesLimit() in google/protobuf/io/coded_stream.h.\n",
      "[libprotobuf WARNING google/protobuf/io/coded_stream.cc:81] The total number of bytes read was 1331138693\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "tcmalloc: large alloc 1800912896 bytes == 0x7f5d74a84000 @  0x7f5e9070c887 0x7f5e11f173ea 0x7f5e11f0a632 0x7f5e120df6d4 0x7f5e11ef638f 0x7f5e1ebca86a 0x7f5e1ec2194a 0x56204c 0x4f88ba 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023 0x6415b2 0x64166a 0x643730 0x62b26e 0x4b4cb0 0x7f5e90307b97 0x5bdf6a\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "INTERENCE TIME: 239.51307100014674 ms\n",
      "TRANSCRIPT:  when these two souls perceived each other they recognized each other as necessary to each other and embraced each other closely\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "PYTHONPATH=/content/DeepLearningExamples/PyTorch/SpeechRecognition/Jasper \n",
    "python ../trt/perf.py \\\n",
    "--ckpt_path ./jasper_fp16.pt --wav=example1.wav \\\n",
    "--model_toml=../configs/jasper10x5dr_nomask.toml \\\n",
    "--make_onnx --onnx_path jasper.onnx \\\n",
    "--engine_path jasper.plan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OFYVKESRdf-y"
   },
   "source": [
    "### Inference from existing TensorRT FP32 plan\n",
    "Inference with an existing plan can be launch with the `--use_existing_engine` flag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 149
    },
    "colab_type": "code",
    "id": "vzQOo6QLODZd",
    "outputId": "54bbe948-02f5-4a3f-eeae-bd4fd8da290c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INTERENCE TIME: 289.92610499994953 ms\n",
      "TRANSCRIPT:  when these two souls perceived each other they recognized each other as necessary to each other and embraced each other closely\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 1331036160 bytes == 0x62440000 @  0x7fd170b6f1e7 0x5a1c5c 0x578954 0x561fca 0x57c961 0x57e6ae 0x4bb666 0x4f858d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023 0x6415b2 0x64166a 0x643730 0x62b26e 0x4b4cb0 0x7fd17076cb97 0x5bdf6a\n",
      "tcmalloc: large alloc 1330552832 bytes == 0x1010f0000 @  0x7fd170b71887 0x7fd0f255dce7 0x7fd0f254d05f 0x7fd0f2364ee3 0x7fd0f236efd8 0x7fd0ff01f82e 0x7fd0ff08694a 0x56204c 0x4f88ba 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023 0x6415b2 0x64166a 0x643730 0x62b26e 0x4b4cb0 0x7fd17076cb97 0x5bdf6a\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "PYTHONPATH=/content/DeepLearningExamples/PyTorch/SpeechRecognition/Jasper \n",
    "python ../trt/perf.py \\\n",
    "--wav=./example1.wav \\\n",
    "--model_toml=../configs/jasper10x5dr_nomask.toml \\\n",
    "--use_existing_engine --engine_path jasper.plan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Pg4X1odgOS1i"
   },
   "source": [
    "## FP16 Inference with TensorRT\n",
    "### Creating TensorRT FP16 execution plan\n",
    "\n",
    "We will next create an FP16 TRT inference plan. \n",
    "\n",
    "To run inference of the input audio file using automatic mixed precision, add the argument `--trt_fp16`. Using automatic mixed precision, the inference time can be reduced efficiently compared to that of using fp32 (building the engine for the first time can take several minutes).\n",
    "\n",
    "**Important Note:** Efficient FP16 inference requires a Volta, Turing or newer generation GPUs. On Google Colab, this normally means a T4 GPU. On the older K80 GPUs, FP16 performance might actually degrade from an FP32 TRT model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 243
    },
    "colab_type": "code",
    "id": "x2n_2cZYdGOg",
    "outputId": "0f02ca21-d5b0-4c16-89db-a46bdf57e438"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INTERENCE TIME: 334.61581900019155 ms\n",
      "TRANSCRIPT:  when these two souls perceived each other they recognized each other as necessary to each other and embraced each other closely\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 1331142656 bytes == 0x15bbec000 @  0x7f1bf09e6887 0x7f1bef2dcbf9 0x7f1bef2ddacb 0x7f1bef2ddb84 0x7f1bef2ddf6c 0x7f1baac2c16f 0x7f1baac2c3f4 0x7f1ba0814411 0x7f1be6b2237d 0x7f1be6857ef4 0x56204c 0x4f88ba 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f7a28 0x4f876d\n",
      "tcmalloc: large alloc 1331142656 bytes == 0x106ce2000 @  0x7f1bf09e41e7 0x5a1c5c 0x7f1be6b226da 0x7f1be6857ef4 0x56204c 0x4f88ba 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4f98c7 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f6128 0x4f7d60 0x4f876d 0x4fa6c0 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023\n",
      "[libprotobuf WARNING google/protobuf/io/coded_stream.cc:604] Reading dangerously large protocol message.  If the message turns out to be larger than 2147483647 bytes, parsing will be halted for security reasons.  To increase the limit (or to disable these warnings), see CodedInputStream::SetTotalBytesLimit() in google/protobuf/io/coded_stream.h.\n",
      "[libprotobuf WARNING google/protobuf/io/coded_stream.cc:81] The total number of bytes read was 1331138693\n",
      "[TensorRT] WARNING: Half2 support requested on hardware without native FP16 support, performance will be negatively affected.\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "tcmalloc: large alloc 1817280512 bytes == 0x7f1ad3ae8000 @  0x7f1bf09e6887 0x7f1b721f13ea 0x7f1b721e4632 0x7f1b723b96d4 0x7f1b721d038f 0x7f1b7eea486a 0x7f1b7eefb94a 0x56204c 0x4f88ba 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023 0x6415b2 0x64166a 0x643730 0x62b26e 0x4b4cb0 0x7f1bf05e1b97 0x5bdf6a\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "PYTHONPATH=/content/DeepLearningExamples/PyTorch/SpeechRecognition/Jasper \n",
    "python ../trt/perf.py \\\n",
    "--ckpt_path ./jasper_fp16.pt --wav=example1.wav \\\n",
    "--model_toml=../configs/jasper10x5dr_nomask.toml \\\n",
    "--make_onnx --onnx_path jasper.onnx \\\n",
    "--engine_path jasper_fp16.plan \\\n",
    "--trt_fp16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "912MBa0BdyTZ"
   },
   "source": [
    "### Inference from existing TensorRT FP16 plan\n",
    "Inference with an existing plan can be launched with the `--use_existing_engine` flag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 131
    },
    "colab_type": "code",
    "id": "C9oVQa_zOV1u",
    "outputId": "2de02886-853c-4deb-9a9c-234a338abec0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INTERENCE TIME: 301.42106899984356 ms\n",
      "TRANSCRIPT:  when these two souls perceived each other they recognized each other as necessary to each other and embraced each other closely\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "tcmalloc: large alloc 1116463104 bytes == 0xb1754000 @  0x7f0533d0e1e7 0x5a1c5c 0x578954 0x561fca 0x57c961 0x57e6ae 0x4bb666 0x4f858d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f7a28 0x4f876d 0x4f98c7 0x4f6128 0x4f9023 0x6415b2 0x64166a 0x643730 0x62b26e 0x4b4cb0 0x7f053390bb97 0x5bdf6a\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n",
      "[TensorRT] WARNING: TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.2\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "PYTHONPATH=/content/DeepLearningExamples/PyTorch/SpeechRecognition/Jasper \n",
    "python ../trt/perf.py \\\n",
    "--wav=./example1.wav \\\n",
    "--model_toml=../configs/jasper10x5dr_nomask.toml \\\n",
    "--use_existing_engine --engine_path jasper_fp16.plan \\\n",
    "--trt_fp16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "g8MxXY5GmTc8"
   },
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we have walked through the complete process of carrying out inference using a pretrained Jasper Pytorch model using NVIDIA TensorRT on Google Colab.\n",
    "### What's next\n",
    "Now that  you are familiar with running Jasper inference with TensorRT using full and automatic mixed precision, you may want to play with your own audio samples.\n",
    "\n",
    "For information on training a Jasper model using your own data, please check out our Github repo: https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/Jasper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Q4_PtLSdMnw3"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "name": "Colab_Jasper_TRT_inference_demo.ipynb",
   "provenance": []
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
