inference_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import time\n",
    "from torchvision.models import *\n",
    "import pandas as pd\n",
    "import os\n",
    "import torchvision\n",
    "from torch2trt import torch2trt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models.segmentation import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "FP32 = True\n",
    "FP16 = True\n",
    "INT8 = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make results\n",
    "os.makedirs(\"results\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelWrapper(torch.nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super(ModelWrapper, self).__init__()\n",
    "        self.model = model\n",
    "    def forward(self, x):\n",
    "        return self.model(x)['out']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def computeTime(model, input_size=[1, 3, 224, 224], device='cuda', FP16=False):\n",
    "    inputs = torch.randn(input_size)\n",
    "    if device == 'cuda':\n",
    "        model = model.cuda()\n",
    "        inputs = inputs.cuda()\n",
    "    if FP16:\n",
    "        model = model.half()\n",
    "        inputs = inputs.half()\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    i = 0\n",
    "    time_spent = []\n",
    "    while i < 200:\n",
    "        start_time = time.time()\n",
    "        with torch.no_grad():\n",
    "            _ = model(inputs)\n",
    "\n",
    "        if device == 'cuda':\n",
    "            torch.cuda.synchronize() # wait for cuda to finish (cuda is asynchronous!)\n",
    "        if i != 0:\n",
    "            time_spent.append(time.time() - start_time)\n",
    "        i += 1\n",
    "    print('Avg execution time (ms): {:.3f}'.format(np.mean(time_spent)))\n",
    "    return np.mean(time_spent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# resnet is enought for now\n",
    "modellist = [\"fcn_resnet50\", \"fcn_resnet101\", \"deeplabv3_resnet50\", \"deeplabv3_resnet101\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model: fcn_resnet50\n",
      "Avg execution time (ms): 0.205\n",
      "Avg execution time (ms): 0.174\n",
      "running fp16 models..\n",
      "Avg execution time (ms): 0.037\n",
      "running int8 models..\n",
      "Avg execution time (ms): 0.022\n",
      "model: fcn_resnet101\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/resnet101-5d3b4d8f.pth\" to /home/ken/.cache/torch/checkpoints/resnet101-5d3b4d8f.pth\n",
      "100.0%\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg execution time (ms): 0.344\n",
      "Avg execution time (ms): 0.290\n",
      "running fp16 models..\n",
      "Avg execution time (ms): 0.057\n",
      "running int8 models..\n",
      "Avg execution time (ms): 0.032\n",
      "model: deeplabv3_resnet50\n",
      "Avg execution time (ms): 0.281\n",
      "Avg execution time (ms): 0.252\n",
      "running fp16 models..\n",
      "Avg execution time (ms): 0.130\n",
      "running int8 models..\n",
      "Avg execution time (ms): 0.097\n",
      "model: deeplabv3_resnet101\n",
      "Avg execution time (ms): 0.426\n",
      "Avg execution time (ms): 0.367\n",
      "running fp16 models..\n",
      "Avg execution time (ms): 0.151\n",
      "running int8 models..\n",
      "Avg execution time (ms): 0.108\n"
     ]
    }
   ],
   "source": [
    "results = []\n",
    "for i, model_name in enumerate(modellist):\n",
    "    runtimes = []\n",
    "\n",
    "    # define model\n",
    "    print(\"model: {}\".format(model_name))\n",
    "    input_size = [1, 3, 512, 512]\n",
    "    mdl = globals()[model_name]\n",
    "    model = mdl().cuda().eval()\n",
    "    # Run raw models\n",
    "    runtimes.append(computeTime(model, input_size=input_size, device=\"cuda\", FP16=False))\n",
    "\n",
    "    if FP32:    \n",
    "        mdl = globals()[model_name]\n",
    "        model = mdl().cuda().eval()\n",
    "        model_w = ModelWrapper(model)\n",
    "        x = torch.zeros(input_size).cuda()\n",
    "\n",
    "        # convert to tensorrt models\n",
    "        model_trt = torch2trt(model_w, [x])\n",
    "\n",
    "        # Run TensorRT models\n",
    "        runtimes.append(computeTime(model_trt, input_size=input_size, device=\"cuda\", FP16=False))\n",
    "    if FP16:\n",
    "        print(\"running fp16 models..\")\n",
    "        # Make FP16 tensorRT models\n",
    "        mdl = globals()[model_name]\n",
    "        model = mdl().eval().half().cuda()\n",
    "        model_w = ModelWrapper(model).half()\n",
    "        x = torch.zeros(input_size).half().cuda()\n",
    "        # convert to tensorrt models\n",
    "        model_trt = torch2trt(model_w, [x], fp16_mode=True)\n",
    "        # Run TensorRT models\n",
    "        runtimes.append(computeTime(model_trt, input_size=input_size, device=\"cuda\", FP16=True))\n",
    "\n",
    "    if INT8:\n",
    "        print(\"running int8 models..\")\n",
    "        # Make INT8 tensorRT models\n",
    "        mdl = globals()[model_name]\n",
    "        model = mdl().eval().half().cuda()\n",
    "        model_w = ModelWrapper(model).half()\n",
    "        x = torch.randn(input_size).half().cuda()\n",
    "        # convert to tensorrt models\n",
    "        model_trt = torch2trt(model_w, [x], fp16_mode=True, int8_mode=True, max_batch_size=1)\n",
    "\n",
    "        runtimes.append(computeTime(model_trt, input_size=input_size, device=\"cuda\", FP16=True))\n",
    "\n",
    "    if i == 0:\n",
    "        df = pd.DataFrame({model_name: runtimes},\n",
    "                         index = [\"Raw\", \"FP32\", \"FP16\", \"INT8\"])\n",
    "    else:\n",
    "        df[model_name] = runtimes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>fcn_resnet50</th>\n",
       "      <th>fcn_resnet101</th>\n",
       "      <th>deeplabv3_resnet50</th>\n",
       "      <th>deeplabv3_resnet101</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Raw</th>\n",
       "      <td>0.205359</td>\n",
       "      <td>0.344307</td>\n",
       "      <td>0.281023</td>\n",
       "      <td>0.425960</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FP32</th>\n",
       "      <td>0.173818</td>\n",
       "      <td>0.290180</td>\n",
       "      <td>0.252314</td>\n",
       "      <td>0.366532</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FP16</th>\n",
       "      <td>0.036635</td>\n",
       "      <td>0.056922</td>\n",
       "      <td>0.129868</td>\n",
       "      <td>0.151195</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>INT8</th>\n",
       "      <td>0.021869</td>\n",
       "      <td>0.032292</td>\n",
       "      <td>0.097351</td>\n",
       "      <td>0.108282</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      fcn_resnet50  fcn_resnet101  deeplabv3_resnet50  deeplabv3_resnet101\n",
       "Raw       0.205359       0.344307            0.281023             0.425960\n",
       "FP32      0.173818       0.290180            0.252314             0.366532\n",
       "FP16      0.036635       0.056922            0.129868             0.151195\n",
       "INT8      0.021869       0.032292            0.097351             0.108282"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.to_csv(\"results/xavier_segmentation.csv\")\n",
    "df"
   ]
  },
  {
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}