notebooks/torch_jit.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `torch.jit`\n",
    "\n",
    "Eager execution is great for development and debugging. but it can be hard to (automatically) optimize the code and deploy it.\n",
    "\n",
    "Now there is`torch.jit` with two flavours:\n",
    "\n",
    "- `torch.jit.trace` does not record control flow.\n",
    "- `torch.jit.script` records control flow and creates an intermediate representation that can be optimized; only supports a subset of Python.\n",
    "\n",
    "Note: don't forget `model.eval()` and `model.train()`.\n",
    "\n",
    "\n",
    "## Ref and More:\n",
    "- https://pytorch.org/docs/stable/jit.html\n",
    "- https://speakerdeck.com/perone/pytorch-under-the-hood\n",
    "- https://lernapparat.de/fast-lstm-pytorch/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Init, helpers, utils, ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:24.971789Z",
     "iopub.status.busy": "2020-11-27T14:31:24.969854Z",
     "iopub.status.idle": "2020-11-27T14:31:25.203359Z",
     "shell.execute_reply": "2020-11-27T14:31:25.202626Z"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.206235Z",
     "iopub.status.busy": "2020-11-27T14:31:25.205811Z",
     "iopub.status.idle": "2020-11-27T14:31:25.478268Z",
     "shell.execute_reply": "2020-11-27T14:31:25.477948Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.481570Z",
     "iopub.status.busy": "2020-11-27T14:31:25.481113Z",
     "iopub.status.idle": "2020-11-27T14:31:25.494485Z",
     "shell.execute_reply": "2020-11-27T14:31:25.494069Z"
    }
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython.core.debugger import set_trace\n",
    "\n",
    "import utils  # little helpers\n",
    "from utils import attr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `torch.jit.trace`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.497670Z",
     "iopub.status.busy": "2020-11-27T14:31:25.497372Z",
     "iopub.status.idle": "2020-11-27T14:31:25.504775Z",
     "shell.execute_reply": "2020-11-27T14:31:25.504504Z"
    }
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    if x.item() < 0:\n",
    "        return torch.tensor(0)\n",
    "    else:\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.507274Z",
     "iopub.status.busy": "2020-11-27T14:31:25.506848Z",
     "iopub.status.idle": "2020-11-27T14:31:25.515959Z",
     "shell.execute_reply": "2020-11-27T14:31:25.515596Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(torch.tensor(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.518286Z",
     "iopub.status.busy": "2020-11-27T14:31:25.517863Z",
     "iopub.status.idle": "2020-11-27T14:31:25.525322Z",
     "shell.execute_reply": "2020-11-27T14:31:25.525004Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(3)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(torch.tensor(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.527966Z",
     "iopub.status.busy": "2020-11-27T14:31:25.527466Z",
     "iopub.status.idle": "2020-11-27T14:31:25.541534Z",
     "shell.execute_reply": "2020-11-27T14:31:25.540887Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-1-5385317fa43b>:2: TracerWarning: Converting a tensor to a Python number might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if x.item() < 0:\n"
     ]
    }
   ],
   "source": [
    "X = torch.tensor(1)\n",
    "traced = torch.jit.trace(f, X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.543997Z",
     "iopub.status.busy": "2020-11-27T14:31:25.543681Z",
     "iopub.status.idle": "2020-11-27T14:31:25.551831Z",
     "shell.execute_reply": "2020-11-27T14:31:25.551525Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.jit.ScriptFunction"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(traced)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.554129Z",
     "iopub.status.busy": "2020-11-27T14:31:25.553805Z",
     "iopub.status.idle": "2020-11-27T14:31:25.562407Z",
     "shell.execute_reply": "2020-11-27T14:31:25.562078Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "traced(torch.tensor(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.564943Z",
     "iopub.status.busy": "2020-11-27T14:31:25.564560Z",
     "iopub.status.idle": "2020-11-27T14:31:25.572611Z",
     "shell.execute_reply": "2020-11-27T14:31:25.572278Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "graph(%0 : Long(requires_grad=0, device=cpu)):\n",
       "  return (%0)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "traced.graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.575102Z",
     "iopub.status.busy": "2020-11-27T14:31:25.574627Z",
     "iopub.status.idle": "2020-11-27T14:31:25.582783Z",
     "shell.execute_reply": "2020-11-27T14:31:25.582382Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-1)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "traced(torch.tensor(-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storing and restoring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.585530Z",
     "iopub.status.busy": "2020-11-27T14:31:25.585143Z",
     "iopub.status.idle": "2020-11-27T14:31:25.606206Z",
     "shell.execute_reply": "2020-11-27T14:31:25.605544Z"
    }
   },
   "outputs": [],
   "source": [
    "traced.save(\"traced.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.608723Z",
     "iopub.status.busy": "2020-11-27T14:31:25.608368Z",
     "iopub.status.idle": "2020-11-27T14:31:25.733329Z",
     "shell.execute_reply": "2020-11-27T14:31:25.731998Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scripted.pt: Zip archive data, at least v?[0] to extract\r\n"
     ]
    }
   ],
   "source": [
    "!file scripted.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.737533Z",
     "iopub.status.busy": "2020-11-27T14:31:25.737240Z",
     "iopub.status.idle": "2020-11-27T14:31:25.745978Z",
     "shell.execute_reply": "2020-11-27T14:31:25.745553Z"
    }
   },
   "outputs": [],
   "source": [
    "g = torch.jit.load(\"traced.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.749203Z",
     "iopub.status.busy": "2020-11-27T14:31:25.748303Z",
     "iopub.status.idle": "2020-11-27T14:31:25.757664Z",
     "shell.execute_reply": "2020-11-27T14:31:25.757985Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g(torch.tensor(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.761208Z",
     "iopub.status.busy": "2020-11-27T14:31:25.760655Z",
     "iopub.status.idle": "2020-11-27T14:31:25.769043Z",
     "shell.execute_reply": "2020-11-27T14:31:25.769360Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-1)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g(torch.tensor(-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `torch.jit.script`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.772118Z",
     "iopub.status.busy": "2020-11-27T14:31:25.771699Z",
     "iopub.status.idle": "2020-11-27T14:31:25.780525Z",
     "shell.execute_reply": "2020-11-27T14:31:25.780141Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bool(torch.tensor(1) < 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.784867Z",
     "iopub.status.busy": "2020-11-27T14:31:25.783207Z",
     "iopub.status.idle": "2020-11-27T14:31:25.792360Z",
     "shell.execute_reply": "2020-11-27T14:31:25.792027Z"
    }
   },
   "outputs": [],
   "source": [
    "@torch.jit.script\n",
    "def f(x):\n",
    "    if bool(x < 0):\n",
    "        result = torch.zeros(1)\n",
    "    else:\n",
    "        result = x\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is `torchscript` which is a only a supset of python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.794784Z",
     "iopub.status.busy": "2020-11-27T14:31:25.794427Z",
     "iopub.status.idle": "2020-11-27T14:31:25.803364Z",
     "shell.execute_reply": "2020-11-27T14:31:25.803087Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(torch.tensor(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.806024Z",
     "iopub.status.busy": "2020-11-27T14:31:25.805284Z",
     "iopub.status.idle": "2020-11-27T14:31:25.813152Z",
     "shell.execute_reply": "2020-11-27T14:31:25.813381Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(torch.tensor(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.815516Z",
     "iopub.status.busy": "2020-11-27T14:31:25.815182Z",
     "iopub.status.idle": "2020-11-27T14:31:25.823160Z",
     "shell.execute_reply": "2020-11-27T14:31:25.822776Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.jit.ScriptFunction"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.825962Z",
     "iopub.status.busy": "2020-11-27T14:31:25.825098Z",
     "iopub.status.idle": "2020-11-27T14:31:25.833882Z",
     "shell.execute_reply": "2020-11-27T14:31:25.833453Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "graph(%x.1 : Tensor):\n",
       "  %8 : None = prim::Constant()\n",
       "  %2 : int = prim::Constant[value=0]() # <ipython-input-1-5b977b5b82b7>:3:16\n",
       "  %5 : int = prim::Constant[value=1]() # <ipython-input-1-5b977b5b82b7>:4:29\n",
       "  %3 : Tensor = aten::lt(%x.1, %2) # <ipython-input-1-5b977b5b82b7>:3:12\n",
       "  %4 : bool = aten::Bool(%3) # <ipython-input-1-5b977b5b82b7>:3:7\n",
       "  %result : Tensor = prim::If(%4) # <ipython-input-1-5b977b5b82b7>:3:4\n",
       "    block0():\n",
       "      %7 : int[] = prim::ListConstruct(%5)\n",
       "      %result.1 : Tensor = aten::zeros(%7, %8, %8, %8, %8) # <ipython-input-1-5b977b5b82b7>:4:17\n",
       "      -> (%result.1)\n",
       "    block1():\n",
       "      -> (%x.1)\n",
       "  return (%result)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f.graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storing and restoring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.836435Z",
     "iopub.status.busy": "2020-11-27T14:31:25.836068Z",
     "iopub.status.idle": "2020-11-27T14:31:25.843708Z",
     "shell.execute_reply": "2020-11-27T14:31:25.843266Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.jit.save(f, \"scripted.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.847133Z",
     "iopub.status.busy": "2020-11-27T14:31:25.846686Z",
     "iopub.status.idle": "2020-11-27T14:31:25.963596Z",
     "shell.execute_reply": "2020-11-27T14:31:25.962539Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scripted.pt: Zip archive data, at least v?[0] to extract\r\n"
     ]
    }
   ],
   "source": [
    "!file scripted.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.968061Z",
     "iopub.status.busy": "2020-11-27T14:31:25.967602Z",
     "iopub.status.idle": "2020-11-27T14:31:25.979424Z",
     "shell.execute_reply": "2020-11-27T14:31:25.978497Z"
    }
   },
   "outputs": [],
   "source": [
    "g = torch.jit.load(\"scripted.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.982405Z",
     "iopub.status.busy": "2020-11-27T14:31:25.982021Z",
     "iopub.status.idle": "2020-11-27T14:31:25.992205Z",
     "shell.execute_reply": "2020-11-27T14:31:25.991823Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g(torch.tensor(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-11-27T14:31:25.995114Z",
     "iopub.status.busy": "2020-11-27T14:31:25.994647Z",
     "iopub.status.idle": "2020-11-27T14:31:26.005153Z",
     "shell.execute_reply": "2020-11-27T14:31:26.004702Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g(torch.tensor(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subclassing `torch.jit.ScriptModule`\n",
    "If you work with `nn.Module` replace it by `torch.jit.ScriptModule` (see [[tutorial]](https://pytorch.org/tutorials/beginner/deploy_seq2seq_hybrid_frontend_tutorial.html) for more).\n",
    "\n",
    "```python\n",
    "class MyModule(torch.jit.ScriptModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # ...\n",
    "        return x\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch and C++\n",
    "\n",
    "PyTorch offers a very nice(!) C++ interface which is very close to Python."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading traced models from C++"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```c++\n",
    "#include <torch/script.h>\n",
    "\n",
    "int main(int(argc, const char* argv[]) {\n",
    "    auto module = torch::jit::load(\"scrpted.pt\");\n",
    "    // data ...\n",
    "    module->forward(data);\n",
    "}\n",
    "```"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}