Colab-torch2trt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colab-torch2trt.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F3khbJqHHBlc"
      },
      "source": [
        "# Colab-torch2trt\n",
        "\n",
        "Original Colab: [here](https://colab.research.google.com/drive/1oe_aflRfCwRehho_8QlD8YFKUQ_I5sA6#scrollTo=tgSrJdHic-Du)\n",
        "\n",
        "torch2trt: [NVIDIA-AI-IOT/torch2trt](https://github.com/NVIDIA-AI-IOT/torch2trt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YwjJ5qneGMri"
      },
      "source": [
        "You need to get 2 files.\n",
        "Currently, the cuda version inside Colab is 11.0, that's why you need to get:\n",
        "\n",
        "```\n",
        "nv-tensorrt-repo-ubuntu1804-cuda11.0-trt7.2.3.4-ga-20210226_1-1_amd64.deb\n",
        "\n",
        "and\n",
        " \n",
        "TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.0.cudnn8.1.tar\\TensorRT-7.2.3.4\\python\\tensorrt-7.2.3.4-cp37-none-linux_x86_64.whl (inside TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.0.cudnn8.1.tar.gz)\n",
        "```\n",
        "\n",
        "You can download these files [here](https://developer.nvidia.com/nvidia-tensorrt-download). Warning: You need an account (which can be created for free).\n",
        "\n",
        "If you want to use other versions, you need to adjust the install script."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "192U9XGDHXSs"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aaazZaDB96v0"
      },
      "source": [
        "!nvcc --version"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l0Os1_hJcnPU",
        "cellView": "form"
      },
      "source": [
        "#@title Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive',force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v1afUfXxB6kw",
        "cellView": "form"
      },
      "source": [
        "#@title install\n",
        "import os\n",
        "os.environ[\"os1\"]=\"ubuntu1804\"\n",
        "os.environ[\"tag\"]= \"cuda11.0-trt7.2.3.4-ga-20210226\" #@param\n",
        "os.environ[\"version\"]= \"7.2.3-1+cuda11.0\" #@param\n",
        "data_path = '/content/drive/MyDrive/tensorrt 11.0/' #@param\n",
        "os.chdir(data_path)\n",
        "!sudo dpkg -i nv-tensorrt-repo-${os1}-${tag}_1-1_amd64.deb\n",
        "!sudo apt-key add /var/nv-tensorrt-repo-${tag}/7fa2af80.pub\n",
        "!sudo apt-get update\n",
        "!sudo apt-get install libnvinfer7=${version} libnvonnxparsers7=${version} libnvparsers7=${version} libnvinfer-plugin7=${version} libnvinfer-dev=${version} libnvonnxparsers-dev=${version} libnvparsers-dev=${version} libnvinfer-plugin-dev=${version} python-libnvinfer=${version} python3-libnvinfer=${version}\n",
        "!sudo apt-mark hold libnvinfer7 libnvonnxparsers7 libnvparsers7 libnvinfer-plugin7 libnvinfer-dev libnvonnxparsers-dev libnvparsers-dev libnvinfer-plugin-dev python-libnvinfer python3-libnvinfer\n",
        "!sudo apt-get install tensorrt=${version}\n",
        "!sudo apt-get install python3-libnvinfer-dev=${version}\n",
        "\n",
        "!git clone https://github.com/NVIDIA-AI-IOT/torch2trt /content/torch2trt\n",
        "%cd /content/torch2trt\n",
        "!python setup.py install"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tgSrJdHic-Du"
      },
      "source": [
        "# IMPORTANT STEP\n",
        "\n",
        "**Restart colab (Runtime > Restart Runtime)**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KOqn4z_O78CO"
      },
      "source": [
        "!pip install \"/content/drive/MyDrive/tensorrt 11.0/tensorrt-7.2.3.4-cp37-none-linux_x86_64.whl\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lwQUiCqhdUhC"
      },
      "source": [
        "#@title Example\n",
        "# testing example from https://github.com/NVIDIA-AI-IOT/torch2trt\n",
        "import torch\n",
        "from torch2trt import torch2trt\n",
        "\n",
        "from torchvision.models.alexnet import alexnet\n",
        "\n",
        "# create some regular pytorch model...\n",
        "model = alexnet(pretrained=True).eval().cuda()\n",
        "\n",
        "# create example data\n",
        "x = torch.ones((1, 3, 224, 224)).cuda()\n",
        "\n",
        "# convert to TensorRT feeding sample data as input\n",
        "model_trt = torch2trt(model, [x])\n",
        "\n",
        "#Using model\n",
        "out = model_trt(x)\n",
        "\n",
        "print(\"Result:\")\n",
        "print(out.shape)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}