1-tvm_pytorch_resnet18_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ResNet(\n",
      "  original_name=ResNet\n",
      "  (conv1): Conv2d(original_name=Conv2d)\n",
      "  (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "  (relu): ReLU(original_name=ReLU)\n",
      "  (maxpool): MaxPool2d(original_name=MaxPool2d)\n",
      "  (layer1): Sequential(\n",
      "    original_name=Sequential\n",
      "    (0): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "    )\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    original_name=Sequential\n",
      "    (0): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (downsample): Sequential(\n",
      "        original_name=Sequential\n",
      "        (0): Conv2d(original_name=Conv2d)\n",
      "        (1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "    )\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    original_name=Sequential\n",
      "    (0): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (downsample): Sequential(\n",
      "        original_name=Sequential\n",
      "        (0): Conv2d(original_name=Conv2d)\n",
      "        (1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "    )\n",
      "  )\n",
      "  (layer4): Sequential(\n",
      "    original_name=Sequential\n",
      "    (0): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (downsample): Sequential(\n",
      "        original_name=Sequential\n",
      "        (0): Conv2d(original_name=Conv2d)\n",
      "        (1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      original_name=BasicBlock\n",
      "      (conv1): Conv2d(original_name=Conv2d)\n",
      "      (bn1): BatchNorm2d(original_name=BatchNorm2d)\n",
      "      (relu): ReLU(original_name=ReLU)\n",
      "      (conv2): Conv2d(original_name=Conv2d)\n",
      "      (bn2): BatchNorm2d(original_name=BatchNorm2d)\n",
      "    )\n",
      "  )\n",
      "  (avgpool): AdaptiveAvgPool2d(original_name=AdaptiveAvgPool2d)\n",
      "  (fc): Linear(original_name=Linear)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import tvm\n",
    "from tvm import relay\n",
    "import numpy as np\n",
    "from tvm.contrib.download import download_testdata\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "from scipy.special import softmax\n",
    "# device = torch.device(\"cpu\")\n",
    "model_name = \"resnet18\"\n",
    "model = getattr(torchvision.models, model_name)(pretrained=True)\n",
    "model = model.eval()\n",
    "\n",
    "# We grab the TorchScripted model via tracing\n",
    "input_shape = [1, 3, 224, 224]\n",
    "input_data = torch.randn(input_shape)\n",
    "scripted_model = torch.jit.trace(model, input_data).eval()\n",
    "print(scripted_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/parker/.tvm_test_data/data/cat.png\n",
      "(224, 224)\n",
      "(1, 3, 224, 224) float32\n"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "\n",
    "img_url = \"https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true\"\n",
    "img_path = download_testdata(img_url, \"cat.png\", module=\"data\")\n",
    "print(img_path)\n",
    "img = Image.open(img_path).resize((224, 224))\n",
    "print(img.size)\n",
    "\n",
    "# Preprocess the image and convert to tensor\n",
    "from torchvision import transforms\n",
    "\n",
    "\n",
    "my_preprocess = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")\n",
    "img = my_preprocess(img)\n",
    "img = np.expand_dims(img, 0)\n",
    "print(img.shape, img.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('input0', (1, 3, 224, 224))]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/parker/workspace/tvm/python/tvm/relay/build_module.py:342: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "######################################################################\n",
    "# Import the graph to Relay\n",
    "# -------------------------\n",
    "# Convert PyTorch graph to Relay graph. The input name can be arbitrary.\n",
    "input_name = \"input0\"\n",
    "shape_list = [(input_name, img.shape)]\n",
    "print(shape_list)\n",
    "mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)\n",
    "\n",
    "######################################################################\n",
    "# Relay Build\n",
    "# -----------\n",
    "# Compile the graph to llvm target with given input specification.\n",
    "target = \"llvm\"\n",
    "target_host = \"llvm\"\n",
    "dev = tvm.cpu(0)\n",
    "with tvm.transform.PassContext(opt_level=7):\n",
    "    lib = relay.build(mod, target=target, target_host=target_host, params=params, mod_name=\"resnet18\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################################################################\n",
    "# Execute the portable graph on TVM\n",
    "# ---------------------------------\n",
    "# Now we can try deploying the compiled model on target.\n",
    "from tvm.contrib import graph_executor\n",
    "\n",
    "m = graph_executor.GraphModule(lib[\"resnet18\"](dev))\n",
    "\n",
    "tvm_time_spent=[]\n",
    "torch_time_spent=[]\n",
    "n_warmup=5\n",
    "n_time=10\n",
    "# tvm_t0 = time.process_time()\n",
    "for i in range(n_warmup+n_time):\n",
    "    dtype = \"float32\"\n",
    "    # Set inputs\n",
    "    m.set_input(input_name, tvm.nd.array(img.astype(dtype)))\n",
    "    tvm_t0 = time.time()\n",
    "    # Execute\n",
    "    m.run()\n",
    "    # Get outputs\n",
    "    tvm_output = m.get_output(0)\n",
    "    tvm_time_spent.append(time.time() - tvm_t0)\n",
    "# tvm_t1 = time.process_time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relay top-1 id: 281, class name: tabby, tabby cat, class probality: 0.36268845200538635\n",
      "Torch top-1 id: 281, class name: tabby, tabby cat, class probality: 0.36268898844718933\n",
      "Relay time(ms): 42.009\n",
      "Torch time(ms): 16.770\n"
     ]
    }
   ],
   "source": [
    "#####################################################################\n",
    "# Look up synset name\n",
    "# -------------------\n",
    "# Look up prediction top 1 index in 1000 class synset.\n",
    "synset_url = \"\".join(\n",
    "    [\n",
    "        \"https://raw.githubusercontent.com/Cadene/\",\n",
    "        \"pretrained-models.pytorch/master/data/\",\n",
    "        \"imagenet_synsets.txt\",\n",
    "    ]\n",
    ")\n",
    "synset_name = \"imagenet_synsets.txt\"\n",
    "synset_path = download_testdata(synset_url, synset_name, module=\"data\")\n",
    "with open(synset_path) as f:\n",
    "    synsets = f.readlines()\n",
    "\n",
    "synsets = [x.strip() for x in synsets]\n",
    "splits = [line.split(\" \") for line in synsets]\n",
    "key_to_classname = {spl[0]: \" \".join(spl[1:]) for spl in splits}\n",
    "\n",
    "class_url = \"\".join(\n",
    "    [\n",
    "        \"https://raw.githubusercontent.com/Cadene/\",\n",
    "        \"pretrained-models.pytorch/master/data/\",\n",
    "        \"imagenet_classes.txt\",\n",
    "    ]\n",
    ")\n",
    "class_name = \"imagenet_classes.txt\"\n",
    "class_path = download_testdata(class_url, class_name, module=\"data\")\n",
    "with open(class_path) as f:\n",
    "    class_id_to_key = f.readlines()\n",
    "\n",
    "class_id_to_key = [x.strip() for x in class_id_to_key]\n",
    "\n",
    "# Get top-1 result for TVM\n",
    "top1_tvm = np.argmax(tvm_output.asnumpy()[0])\n",
    "tvm_class_key = class_id_to_key[top1_tvm]\n",
    "\n",
    "# Convert input to PyTorch variable and get PyTorch result for comparison\n",
    "# torch_t0 = time.process_time()\n",
    "# torch.set_num_threads(1)\n",
    "for i in range(n_warmup+n_time):\n",
    "    with torch.no_grad():\n",
    "        torch_img = torch.from_numpy(img)\n",
    "        torch_t0 = time.time()\n",
    "        output = model(torch_img)\n",
    "        torch_time_spent.append(time.time() - torch_t0)\n",
    "        # Get top-1 result for PyTorch\n",
    "        top1_torch = np.argmax(output.numpy())\n",
    "        torch_class_key = class_id_to_key[top1_torch]\n",
    "# torch_t1 = time.process_time()\n",
    "\n",
    "# tvm_time = tvm_t1 - tvm_t0\n",
    "# torch_time = torch_t1 - torch_t0\n",
    "tvm_time = np.mean(tvm_time_spent[n_warmup:]) * 1000\n",
    "torch_time = np.mean(torch_time_spent[n_warmup:]) * 1000\n",
    "tvm_output_prob = softmax(tvm_output.asnumpy())\n",
    "output_prob = softmax(output.numpy())\n",
    "print(\"Relay top-1 id: {}, class name: {}, class probality: {}\".format(top1_tvm, key_to_classname[tvm_class_key], tvm_output_prob[0][top1_tvm]))\n",
    "print(\"Torch top-1 id: {}, class name: {}, class probality: {}\".format(top1_torch, key_to_classname[torch_class_key], output_prob[0][top1_torch]))\n",
    "print('Relay time(ms): {:.3f}'.format(tvm_time))\n",
    "print('Torch time(ms): {:.3f}'.format(torch_time))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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.8.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}