06_Optimizing_runtime_with_TorchScript.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "82f9313a-9a67-41f7-ae47-1a464e5aebae",
   "metadata": {},
   "source": [
    "PyTorch is optimized to perform operations on large tensors. Doing many operations on small tensors is quite inefficient in PyTorch. So, whenever possible you should rewrite your computations in batch form to reduce overhead and improve performance. If there's no way you can manually batch your operations, using TorchScript may improve your code's performance. TorchScript is simply a subset of Python functions that are recognized by PyTorch. PyTorch can automatically optimize your TorchScript code using its just in time (jit) compiler and reduce some overheads.\n",
    "\n",
    "Let's look at an example. A very common operation in ML applications is \"batch gather\". This operation can simply written as `output[i] = input[i, index[i]]`. \n",
    "\n",
    "Suppose you have a tensor `input` of shape `(3, 4)`:\n",
    "\n",
    "```\n",
    "input = [[ 1,  2,  3,  4],\n",
    "         [ 5,  6,  7,  8],\n",
    "         [ 9, 10, 11, 12]]\n",
    "```\n",
    "\n",
    "And you have an index tensor `index` of shape `(3, 2)`:\n",
    "\n",
    "```\n",
    "index = [[0, 2],\n",
    "         [1, 3],\n",
    "         [2, 0]]\n",
    "```\n",
    "\n",
    "Now, let's say you want to gather elements from `input` according to the indices specified in `index`. \n",
    "\n",
    "For example, for the first row of `input`, you want to gather elements at indices `0` and `2`, for the second row, you want to gather elements at indices `1` and `3`, and so on.\n",
    "\n",
    "Here's how you can use `torch.gather` to achieve this:\n",
    "\n",
    "```python\n",
    "import torch\n",
    "\n",
    "input = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])\n",
    "index = torch.tensor([[0, 2], [1, 3], [2, 0]])\n",
    "\n",
    "output = torch.gather(input, 1, index)\n",
    "print(output)\n",
    "```\n",
    "\n",
    "This will give you the following output:\n",
    "\n",
    "```\n",
    "tensor([[ 1,  3],\n",
    "        [ 6,  8],\n",
    "        [11,  9]])\n",
    "```\n",
    "\n",
    "Explanation:\n",
    "- For the first row of `input`, `torch.gather` selects elements at indices `0` and `2`, which are `1` and `3`.\n",
    "- For the second row of `input`, it selects elements at indices `1` and `3`, which are `6` and `8`.\n",
    "- For the third row of `input`, it selects elements at indices `2` and `0`, which are `11` and `9`.\n",
    "\n",
    "`torch.gather` allows you to gather specific elements from a tensor along a specified dimension according to the indices provided in another tensor.\n",
    "\n",
    "\n",
    "In essence, batch gather enables you to gather elements from multiple batches of tensors based on corresponding indices. This can be useful in various scenarios, such as when you have a batch of sequences (e.g., sequences of words in natural language processing tasks) and you want to gather specific elements from each sequence according to indices provided for each batch.\n",
    "\n",
    "This can be simply implemented in PyTorch as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c2f5a677-d029-481f-8617-57d8dcd2e7a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "def batch_gather(tensor, indices):\n",
    "    output = []\n",
    "    for i in range(tensor.size(0)):\n",
    "        output += [tensor[i][indices[i]]]\n",
    "    return torch.stack(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "748f25c2-15d5-418a-b52a-8f64b403deb9",
   "metadata": {},
   "source": [
    "To implement the same function using TorchScript simply use the `torch.jit.script` decorator:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "85288901-cff2-4800-be17-3b0e6f47fbef",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.jit.script\n",
    "def batch_gather_jit(tensor, indices):\n",
    "    output = []\n",
    "    for i in range(tensor.size(0)):\n",
    "        output += [tensor[i][indices[i]]]\n",
    "    return torch.stack(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "425c2819-12e6-43e0-a751-42ce3f526800",
   "metadata": {},
   "source": [
    "On my tests this is about 10% faster.\r\n",
    "\r\n",
    "But nothing beats manually batching your operations. A vectorized implementation in my tests is 100 times faster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e836ebb2-d286-4a89-981f-056638ad98a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_gather_vec(tensor, indices):\n",
    "    shape = list(tensor.shape)\n",
    "    flat_first = torch.reshape(\n",
    "        tensor, [shape[0] * shape[1]] + shape[2:])\n",
    "    offset = torch.reshape(\n",
    "        torch.arange(shape[0]).cuda() * shape[1],\n",
    "        [shape[0]] + [1] * (len(indices.shape) - 1))\n",
    "    output = flat_first[indices + offset]\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49fb9eb3-aaf6-4e92-8eba-cd27d30a89c0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}