PyTorch/Inference/Gaudi_inference_ex2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c051f912",
   "metadata": {},
   "source": [
    "Copyright (c) 2022 Habana Labs, Ltd. an Intel Company.\n",
    "All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8e01d0c",
   "metadata": {},
   "source": [
    "# Licensed under the Apache License, Version 2.0 (the “License”);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf0eddf",
   "metadata": {},
   "source": [
    "you may not use this file except in compliance with the License. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f5b6012",
   "metadata": {},
   "source": [
    "# Inference on Gaudi - Example2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "561e4869",
   "metadata": {},
   "source": [
    "This notebook is used as an example to show inference on the Gaudi Accelerator. This is using a simple model and the MNIST dataset.\n",
    "\n",
    "This tutorial will show how to infer an MNIST model with HPU GRAPH using HPU Graph APIs and Stream APIs. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1d94b8",
   "metadata": {},
   "source": [
    "Download pretrained model checkpoints from vault"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c83a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://vault.habana.ai/artifactory/misc/inference/mnist/mnist-epoch_20.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e032945e",
   "metadata": {},
   "source": [
    "Import all neccessary dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7f64868a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import time\n",
    "import habana_frameworks.torch as ht\n",
    "import habana_frameworks.torch.core as htcore\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms, datasets\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0943d83d",
   "metadata": {},
   "source": [
    "Define a simple Net model for MNIST.\n",
    "Create the model, and load the pre-trained checkpoint.\n",
    "Optimize the model for eval, and move the model to Gaudi(hpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "50f2a4d9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1   = nn.Linear(784, 256)\n",
    "        self.fc2   = nn.Linear(256, 64)\n",
    "        self.fc3   = nn.Linear(64, 10)\n",
    "    def forward(self, x):\n",
    "        out = x.view(-1,28*28)\n",
    "        out = F.relu(self.fc1(out))\n",
    "        out = F.relu(self.fc2(out))\n",
    "        out = self.fc3(out)\n",
    "        out = F.log_softmax(out, dim=1)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6819209c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "=============================HABANA PT BRIDGE CONFIGURATION =========================== \n",
      " PT_HPU_LAZY_MODE = 1\n",
      " PT_HPU_LAZY_EAGER_OPTIM_CACHE = 1\n",
      " PT_HPU_ENABLE_COMPILE_THREAD = 0\n",
      " PT_HPU_ENABLE_EXECUTION_THREAD = 1\n",
      " PT_HPU_ENABLE_LAZY_EAGER_EXECUTION_THREAD = 1\n",
      " PT_ENABLE_INTER_HOST_CACHING = 0\n",
      " PT_ENABLE_INFERENCE_MODE = 1\n",
      " PT_ENABLE_HABANA_CACHING = 1\n",
      " PT_HPU_MAX_RECIPE_SUBMISSION_LIMIT = 0\n",
      " PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
      " PT_HPU_MAX_COMPOUND_OP_SIZE_SS = 10\n",
      " PT_HPU_ENABLE_STAGE_SUBMISSION = 1\n",
      " PT_HPU_STAGE_SUBMISSION_MODE = 2\n",
      " PT_HPU_PGM_ENABLE_CACHE = 1\n",
      " PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
      " PT_HCCL_SLICE_SIZE_MB = 16\n",
      " PT_HCCL_MEMORY_ALLOWANCE_MB = 0\n",
      " PT_HPU_INITIAL_WORKSPACE_SIZE = 0\n",
      " PT_HABANA_POOL_SIZE = 24\n",
      " PT_HPU_POOL_STRATEGY = 5\n",
      " PT_HPU_POOL_LOG_FRAGMENTATION_INFO = 0\n",
      " PT_ENABLE_MEMORY_DEFRAGMENTATION = 1\n",
      " PT_ENABLE_DEFRAGMENTATION_INFO = 0\n",
      " PT_HPU_ENABLE_SYNAPSE_LAYOUT_HANDLING = 1\n",
      " PT_HPU_ENABLE_SYNAPSE_OUTPUT_PERMUTE = 1\n",
      " PT_HPU_ENABLE_VALID_DATA_RANGE_CHECK = 1\n",
      " PT_HPU_FORCE_USE_DEFAULT_STREAM = 0\n",
      " PT_RECIPE_CACHE_PATH = \n",
      " PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
      " PT_HPU_DYNAMIC_MIN_POLICY_ORDER = 4,5,3,1\n",
      " PT_HPU_DYNAMIC_MAX_POLICY_ORDER = 2,4,5,3,1\n",
      " PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
      "=============================SYSTEM CONFIGURATION ========================================= \n",
      "Num CPU Cores = 96\n",
      "CPU RAM = 784300908 KB \n",
      "============================================================================================ \n"
     ]
    }
   ],
   "source": [
    "model = Net()\n",
    "checkpoint = torch.load('mnist-epoch_20.pth')\n",
    "model.load_state_dict(checkpoint)\n",
    "model = model.eval()\n",
    "\n",
    "model = model.to(\"hpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad5d38f6",
   "metadata": {},
   "source": [
    "Create a MNIST datasets for evaulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "501d0838",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))])\n",
    "\n",
    "data_path = './data'\n",
    "test_kwargs = {'batch_size': 32}\n",
    "dataset1 = datasets.MNIST(data_path, train=False, download=True, transform=transform)\n",
    "test_loader = torch.utils.data.DataLoader(dataset1,**test_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27596215",
   "metadata": {},
   "source": [
    "Create HPUGraph and HPUStream\n",
    "\n",
    "The HPUGraph API provides a performance optimization technique to reduce PyTorch host overhead. This is done by capturing the PyTorch execution on a stream for the first iteration and replaying that in subsequent ones. The replay avoids the PyTorch overhead of accumulating the ops in the model and makes the execution device bound.\n",
    "\n",
    "For further details on Stream APIs and HPU Graph APIs, refer to HPU Graph APIs and Stream APIs in the documentation.\n",
    "\n",
    "The example below shows the capture and replay of HPU Graphs using the following functions:\n",
    "\n",
    "capture_begin() Begins capturing HPU work on the current stream.\n",
    "capture_end() Ends capturing HPU work on the current stream.\n",
    "replay() Replays the HPU work captured by this graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "530e4830",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "g = ht.hpu.HPUGraph()\n",
    "s = ht.hpu.Stream()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1444ac15",
   "metadata": {},
   "source": [
    "Do warm run and capture the HPUGraph \n",
    "\n",
    "Here we are using capture and replay of HPU Graphs:\n",
    "capture_begin() Begins capturing HPU work on the current stream.\n",
    "capture_end() Ends capturing HPU work on the current stream.\n",
    "replay() Replays the HPU work captured by this graph.\n",
    "\n",
    "The initial warmup step is run before replaying the graph in the subsequent infernce pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b474ef1b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "static_input = None\n",
    "output = None\n",
    "\n",
    "#WARMRUN \n",
    "static_input = torch.randn(32, 1, 28, 28, device='hpu')\n",
    "with ht.hpu.stream(s):\n",
    "    g.capture_begin()\n",
    "    output = model(static_input)\n",
    "    g.capture_end()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d72b587",
   "metadata": {},
   "source": [
    "Run inference by replaying the graph. \n",
    "\n",
    "Here, we need to copy the input data to the input placeholder(static_input), and then replay. Output will be saved in the output placeholder(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "647f2da3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 94.36%\n"
     ]
    }
   ],
   "source": [
    "correct = 0 \n",
    "for batch_idx, (data, label) in enumerate(test_loader):  \n",
    "    static_input.copy_(data)\n",
    "    g.replay()\n",
    "    if output.shape[0] != label.shape[0]:\n",
    "        output = output[:label.shape[0]]\n",
    "    correct += output.max(1)[1].eq(label).sum()\n",
    "\n",
    "print('Accuracy: {:.2f}%'.format(100. * correct / (len(test_loader) * 32)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "444ba661-d480-446d-b719-220a476261f9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "quit()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}