transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "96cf5d41-d2a8-47a8-ab74-c1bc5b546062",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import copy\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from torch.nn import (\n",
    "    Module,\n",
    "    ModuleList,\n",
    "    Sequential,\n",
    "    Parameter,\n",
    "    Linear, \n",
    "    Dropout,\n",
    "    LayerNorm,\n",
    "    Softmax,\n",
    ")\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from einops import rearrange as re\n",
    "from opt_einsum import contract as einsum\n",
    "%load_ext line_profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "52a0df6b-3e48-413a-8b22-30039a880740",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "torch.use_deterministic_algorithms(True)\n",
    "random.seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "62c61cb0-2b8d-4d2d-88a6-9ef6bec3ca55",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_printoptions(precision=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "217668ed-d993-4f1f-808a-7b1b1e467b3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def heat(x):\n",
    "    df = pd.DataFrame(x.detach().numpy())\n",
    "    return df.style.background_gradient(cmap='Blues')  # .format('{:.0f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39030dd9-81af-4c70-8567-1994e592bf35",
   "metadata": {},
   "source": [
    "[Attention is All You Need](https://arxiv.org/abs/1706.03762?context=cs)  \n",
    "[Formal Algorithms for Transformers](https://arxiv.org/abs/2207.09238)  \n",
    "[Transformer Language Model Mathematical Definition](https://www.apronus.com/math/transformer-language-model-definition)  \n",
    "[AI Explained - 3D viz of transformer structure](https://youtu.be/-9vVhYEXeyQ?t=456)  \n",
    "[Bloem Transformer Implementation source](https://github.com/pbloem/former/blob/master/former/modules.py)  \n",
    "[The Annotated Transformer](http://nlp.seas.harvard.edu/annotated-transformer/#encoder-and-decoder-stacks)  \n",
    "[Pytorch multi_head_attention_forward source](https://github.com/pytorch/pytorch/blob/dcf51885618e7d1d9aa6e628f3354f67ad82b446/torch/nn/functional.py#L4917)   \n",
    "[Writing a better code with pytorch and einops](http://einops.rocks/pytorch-examples.html)  \n",
    "[R-Drop: Regularized Dropout for Neural Networks](https://arxiv.org/pdf/2106.14448v2.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ffafc46e-0b8f-41b3-be52-c0856ec506b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, d_model=512, n_blocks=6, vocab=30_000):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(vocab, d_model)\n",
    "        self.encoder = Stack(EncoderBlock, n_blocks)\n",
    "        self.decoder = Stack(DecoderBlock, n_blocks)\n",
    "        self.head = Sequential(Linear(d_model, vocab), Softmax(dim=-1))\n",
    "       \n",
    "    def forward(self, src, tgt, src_mask, tgt_mask):\n",
    "        ctx = self.encoder(src, mask=src_mask)\n",
    "        return self.head(self.decoder(tgt, ctx=ctx, mask=tgt_mask, ctx_mask=tgt_mask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4874cf2f-fc73-4dfe-804c-c36c15789d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Stack(nn.Sequential):\n",
    "    def __init__(self, Layer, n):\n",
    "        stack = [copy.deepcopy(Layer()) for l in range(n)]\n",
    "        super().__init__(*stack)\n",
    "        \n",
    "    def forward(self, x, *args, **kwargs):\n",
    "        for module in self:\n",
    "            x = module(x, *args, **kwargs)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "08d23126-22ce-4424-b818-26bab1b2a085",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.mha = MultiHeadAttention()\n",
    "        self.ff = FeedForward()\n",
    "        \n",
    "    def forward(self, x, *, mask=None):\n",
    "        return self.ff(self.mha(x, mask=mask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "77fac8b5-9f09-4f55-a830-01ec26064552",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.masked_mha = MultiHeadAttention()\n",
    "        self.mha = MultiHeadAttention()\n",
    "        self.ff = FeedForward()\n",
    "        \n",
    "    def forward(self, x, *, ctx=None, mask=None, ctx_mask=None):\n",
    "        x = self.masked_mha(x, mask=mask)\n",
    "        return self.ff(self.mha(x, ctx, ctx, ctx_mask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ac798568-9abf-4811-9529-6f68b5cebae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, d_model=512, d_ff=2048):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(d_model, d_ff)\n",
    "        self.fc2 = nn.Linear(d_ff, d_model)\n",
    "        self.norm = nn.LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.norm(x + self._ff(x))\n",
    "    \n",
    "    def _ff(self, x):\n",
    "        return self.fc2(F.relu(self.fc1(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b5594a1c-9d36-4a90-abc1-4f830c893704",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, n_heads: int = 8, emb_dim: int = 512):\n",
    "        super().__init__()\n",
    "        assert emb_dim % n_heads == 0  # d_h = d_q = d_k = d_v = emb_dim // n_heads\n",
    "        self.proj_qkv = ModuleList([Linear(emb_dim, emb_dim, bias=False) for _ in range(3)])\n",
    "        self.proj_o = Linear(emb_dim, emb_dim, bias=False)\n",
    "        self.norm = nn.LayerNorm(emb_dim)\n",
    "        self.h, self.d = n_heads, emb_dim\n",
    "\n",
    "    def forward(self, q, k=None, v=None, mask=None):\n",
    "        \"\"\" q, k, v: (batch, seq_len, emb_dim) mask: (seq_len, seq_len) \"\"\"\n",
    "        if k is None and v is None:\n",
    "            k = v = q\n",
    "        if mask is None:\n",
    "            mask = torch.zeros((q.shape[-2]))\n",
    "        return self.norm(q + self._mha(q, k, v, mask))\n",
    "    \n",
    "    def _mha(self, q, k, v, mask):\n",
    "        q, k, v = (proj(x) for x, proj in zip((q, k, v), self.proj_qkv))\n",
    "        q, k, v = (re(x, \"b l (h d) -> b h l d\", h=self.h) for x in (q, k, v))\n",
    "        attn = einsum(\"...ij,...kj->...ik\", q, k)\n",
    "        attn = mask + torch.einsum(\"...ij,...kj->...ik\", q, k)\n",
    "        attn = F.softmax(attn / q.shape[-1] ** (1/2), dim=-1)\n",
    "        out = einsum(\"...ij,...jk->...ik\", attn, v)\n",
    "        out = re(out, \"b h n d -> b n (h d)\")\n",
    "        out = self.proj_o(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6132f4e3-a172-46d5-9e1a-f490b0c578c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MHAInference(MultiHeadAttention):\n",
    "    def __init__(self, n_heads: int = 8, emb_dim: int = 512):\n",
    "        super().__init__(n_heads, emb_dim)\n",
    "        self.attn = torch.empty((1000 * 8, 256, 256))\n",
    "        self.out = torch.empty((1000 * 8, 256, 64))\n",
    "\n",
    "    def _mha(self, q, k, v, mask):\n",
    "        with torch.no_grad():\n",
    "            q, k, v = (proj(x) for x, proj in zip((q, k, v), self.proj_qkv))\n",
    "            q, k, v = (re(x, \"b l (h d_h) -> (b h) l d_h\", h=self.h) for x in (q, k, v))\n",
    "            torch.bmm(q, re(k, \"bh l d_h -> bh d_h l\"), out=self.attn)\n",
    "            self.attn += mask\n",
    "            self.attn /= q.shape[-1] ** (1/2)\n",
    "            self.attn = F.softmax(self.attn, dim=-1)\n",
    "            torch.bmm(self.attn, v, out=self.out)\n",
    "            return self.proj_o(re(self.out, \"(b h) l d_h -> b l (h d_h)\", h=self.h))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af15b557-e0de-4492-8d92-ee6a1980b05f",
   "metadata": {},
   "source": [
    "# Layer Checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "39a8c7db-7f3b-4f6d-a9b8-55dd289adc31",
   "metadata": {},
   "outputs": [],
   "source": [
    "b, l, d = 1, 256, 512\n",
    "sh = (b, l, d)\n",
    "v = 30_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "0907fdae-0c1c-4fa8-b998-0395b4a26f11",
   "metadata": {},
   "outputs": [],
   "source": [
    "emb = torch.rand(sh)\n",
    "mask = torch.triu(torch.full((l, l), -torch.inf), diagonal=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "acc8214e-63db-436b-8050-291bb442abb6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 54 ms, sys: 23.5 ms, total: 77.5 ms\n",
      "Wall time: 21.6 ms\n"
     ]
    }
   ],
   "source": [
    "%time assert MultiHeadAttention()(emb).shape == sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e889d135-a7fc-4246-9432-201c78e673ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 67.8 ms, sys: 7.96 ms, total: 75.8 ms\n",
      "Wall time: 20.4 ms\n"
     ]
    }
   ],
   "source": [
    "%time assert FeedForward()(emb).shape == sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "90d5c7c1-4095-4f84-aabc-29b67d1db6e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 123 ms, sys: 30.3 ms, total: 154 ms\n",
      "Wall time: 36.3 ms\n"
     ]
    }
   ],
   "source": [
    "%time assert EncoderBlock()(emb).shape == sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c21fe8f9-9356-4851-a442-f9e8d7989755",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 172 ms, sys: 32.4 ms, total: 204 ms\n",
      "Wall time: 51.6 ms\n"
     ]
    }
   ],
   "source": [
    "%time assert DecoderBlock()(emb).shape == sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "e2c022c3-d840-427b-91ed-b68881cb1601",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 58.5 ms, sys: 9.98 ms, total: 68.5 ms\n",
      "Wall time: 14.4 ms\n"
     ]
    }
   ],
   "source": [
    "%time assert MultiHeadAttention()(emb).shape == sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "45502bc0-6551-44ee-85c5-3c2f860c6c8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.49 s, sys: 229 ms, total: 3.72 s\n",
      "Wall time: 486 ms\n"
     ]
    }
   ],
   "source": [
    "%time assert Transformer()(emb, emb, mask, mask).shape == (b, l, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a23d64f-7625-4e9e-80cc-5d775774eaef",
   "metadata": {},
   "source": [
    "# Compare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b5b1fe18-6173-43cb-aa3e-e50b069287f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "emb = torch.rand((1_000, 256, 512))\n",
    "mask = torch.triu(torch.full((256, 256), -torch.inf), diagonal=1)\n",
    "mha = MultiHeadAttention()\n",
    "mha.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "336f18c0-0c62-405d-ae78-6729e4107667",
   "metadata": {},
   "source": [
    "#### _mha vs _mha_infer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5ad3e64c-aeee-44a8-942f-a06e9e1ca918",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.83 s ± 69.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "_ = mha._mha(emb, emb, emb, mask=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "284e0116-0df5-4d59-a314-dc5bd68cd7e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "mha_infer = MHAInference()\n",
    "mha_infer.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "bca96857-4056-4c8f-a143-e9fc68221a3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.32 s ± 27.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "_ = mha_infer(emb, emb, emb, mask=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a32035f1-0a9c-4802-96d9-55dd8e113dd6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Timer unit: 0.001 s\n",
       "\n",
       "Total time: 2.88578 s\n",
       "File: /var/folders/5y/b092b3m96yb8nglxy9dzqbnr0000gn/T/ipykernel_51150/1451311884.py\n",
       "Function: _mha at line 68\n",
       "\n",
       "Line #      Hits         Time  Per Hit   % Time  Line Contents\n",
       "==============================================================\n",
       "    68                                               def _mha(self, q, k, v, mask):\n",
       "    69         1        633.7    633.7     22.0          q, k, v = (proj(x) for x, proj in zip((q, k, v), self.proj_qkv))\n",
       "    70         1          0.1      0.1      0.0          q, k, v = (re(x, \"b l (h d) -> b h l d\", h=self.h) for x in (q, k, v))\n",
       "    71         1        355.6    355.6     12.3          attn = einsum(\"...ij,...kj->...ik\", q, k)\n",
       "    72         1        840.5    840.5     29.1          attn = mask + torch.einsum(\"...ij,...kj->...ik\", q, k)\n",
       "    73         1        633.2    633.2     21.9          attn = F.softmax(attn / q.shape[-1] ** (1/2), dim=-1)\n",
       "    74         1        200.7    200.7      7.0          out = einsum(\"...ij,...jk->...ik\", attn, v)\n",
       "    75         1          0.1      0.1      0.0          out = re(out, \"b h n d -> b n (h d)\")\n",
       "    76         1        221.9    221.9      7.7          out = self.proj_o(out)\n",
       "    77         1          0.0      0.0      0.0          return out"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%lprun -u 0.001 -f mha._mha mha(emb, emb, emb, mask=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6ee3cce2-44d2-4ca7-b677-a11b9c293f6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Timer unit: 0.001 s\n",
       "\n",
       "Total time: 0 s\n",
       "File: /var/folders/5y/b092b3m96yb8nglxy9dzqbnr0000gn/T/ipykernel_51150/1451311884.py\n",
       "Function: _mha at line 68\n",
       "\n",
       "Line #      Hits         Time  Per Hit   % Time  Line Contents\n",
       "==============================================================\n",
       "    68                                               def _mha(self, q, k, v, mask):\n",
       "    69                                                   q, k, v = (proj(x) for x, proj in zip((q, k, v), self.proj_qkv))\n",
       "    70                                                   q, k, v = (re(x, \"b l (h d) -> b h l d\", h=self.h) for x in (q, k, v))\n",
       "    71                                                   attn = einsum(\"...ij,...kj->...ik\", q, k)\n",
       "    72                                                   attn = mask + torch.einsum(\"...ij,...kj->...ik\", q, k)\n",
       "    73                                                   attn = F.softmax(attn / q.shape[-1] ** (1/2), dim=-1)\n",
       "    74                                                   out = einsum(\"...ij,...jk->...ik\", attn, v)\n",
       "    75                                                   out = re(out, \"b h n d -> b n (h d)\")\n",
       "    76                                                   out = self.proj_o(out)\n",
       "    77                                                   return out"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%lprun -u 0.001 -f mha_infer._mha mha_infer(emb, emb, emb, mask=mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6f45895-2166-41f1-84b1-d4efd7dbc988",
   "metadata": {},
   "source": [
    "#### Pytorch [MultiheadAttention.forward](https://github.com/pytorch/pytorch/blob/bbe8d019f280478dc3b143f6988e3e5668499f28/torch/nn/modules/activation.py#L1010) local\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f4fbd57a-31b7-4922-984a-ea7c6bad2a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "th_emb = re(emb, 'b l d -> l b d')\n",
    "bsz, tgt_len, embed_dim = emb.shape\n",
    "num_heads, head_dim = 8, 64\n",
    "in_proj_weight = nn.Parameter(torch.vstack([l.weight for l in mha.proj_qkv]))\n",
    "out_proj_weight = mha.proj_o.weight\n",
    "th_proj_o = Linear(512, 512, bias=False)\n",
    "th_proj_o.weight = out_proj_weight\n",
    "\n",
    "def th_mha():\n",
    "    th_q, th_k, th_v = F._in_projection_packed(th_emb, th_emb, th_emb, in_proj_weight, None)\n",
    "    th_q = th_q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
    "    th_k = th_k.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
    "    th_v = th_v.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
    "    th_out, th_attn = F._scaled_dot_product_attention(th_q, th_k, th_v, mask, 0.0)\n",
    "    th_out = th_out.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)\n",
    "    th_out = th_proj_o(th_out)\n",
    "    th_out = th_out.view(tgt_len, bsz, th_out.size(1))\n",
    "    th_out = mha.norm(th_emb + th_out)\n",
    "    return re(th_out, 'l b d -> b l d')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d78bbfab-22f4-414f-95d4-7a468bb2a046",
   "metadata": {},
   "source": [
    "#### vs. Pytorch F.multi_head_attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "71e20efa-d1a7-47ca-9eee-d72a98bc1207",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.13 s ± 92.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "mha_out = mha(emb, mask=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ba9bc7d4-0c92-43d6-b5bc-e78ee503c20e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.67 s ± 43.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "th_out = th_mha()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b08b0f47-adff-4cbb-8b71-c3eb35d68d17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Timer unit: 0.001 s\n",
       "\n",
       "Total time: 2.59382 s\n",
       "File: /var/folders/5y/b092b3m96yb8nglxy9dzqbnr0000gn/T/ipykernel_51150/486783080.py\n",
       "Function: th_mha at line 9\n",
       "\n",
       "Line #      Hits         Time  Per Hit   % Time  Line Contents\n",
       "==============================================================\n",
       "     9                                           def th_mha():\n",
       "    10         1        821.6    821.6     31.7      th_q, th_k, th_v = F._in_projection_packed(th_emb, th_emb, th_emb, in_proj_weight, None)\n",
       "    11         1         52.8     52.8      2.0      th_q = th_q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
       "    12         1         55.9     55.9      2.2      th_k = th_k.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
       "    13         1        126.6    126.6      4.9      th_v = th_v.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
       "    14         1        989.2    989.2     38.1      th_out, th_attn = F._scaled_dot_product_attention(th_q, th_k, th_v, mask, 0.0)\n",
       "    15         1         87.4     87.4      3.4      th_out = th_out.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)\n",
       "    16         1        223.0    223.0      8.6      th_out = th_proj_o(th_out)\n",
       "    17         1          0.0      0.0      0.0      th_out = th_out.view(tgt_len, bsz, th_out.size(1))\n",
       "    18         1        237.2    237.2      9.1      th_out = mha.norm(th_emb + th_out)\n",
       "    19         1          0.1      0.1      0.0      return re(th_out, 'l b d -> b l d')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%lprun -u 0.001 -f th_mha th_out = th_mha()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7fa59170-ae37-4b41-a410-7887ccebb1e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mha_out = mha(emb, mask=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "468aaf29-8fb0-468b-94df-5e10aea829c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(mha_out, th_out, atol=1e-6, equal_nan=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29a3a853-934e-4ebd-b5f6-f46c11284d5a",
   "metadata": {},
   "source": [
    "## ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ae9e9cbf-1627-4abe-8777-b91c50fd3ec0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/davidwagner/anaconda3/envs/ap/lib/python3.7/site-packages/einops/einops.py:202: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n",
      "  inferred_length: int = length // known_product\n",
      "/Users/davidwagner/anaconda3/envs/ap/lib/python3.7/site-packages/opt_einsum/contract.py:231: TracerWarning: Converting a tensor to a Python integer 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",
      "  dim = int(sh[cnum])\n",
      "/Users/davidwagner/anaconda3/envs/ap/lib/python3.7/site-packages/opt_einsum/parser.py:155: TracerWarning: Converting a tensor to a Python boolean 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",
      "  max(shape[loc] for shape, loc in zip(shapes, [x.find(c) for x in inputs]) if loc >= 0) for c in output)\n"
     ]
    }
   ],
   "source": [
    "torch.onnx.export(mha, emb.detach(), \"mha.onnx\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "004a048a-3990-4566-824c-8427393163d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "onnx_model = onnx.load(\"mha.onnx\")\n",
    "onnx.checker.check_model(onnx_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9e6272d4-0494-4327-b323-e5329af22826",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "graph torch_jit (\n",
      "  %onnx::MatMul_0[FLOAT, 1000x256x512]\n",
      ") initializers (\n",
      "  %norm.weight[FLOAT, 512]\n",
      "  %norm.bias[FLOAT, 512]\n",
      "  %onnx::MatMul_164[FLOAT, 512x512]\n",
      "  %onnx::MatMul_165[FLOAT, 512x512]\n",
      "  %onnx::MatMul_166[FLOAT, 512x512]\n",
      "  %onnx::Reshape_171[INT64, 4]\n",
      "  %onnx::Reshape_176[INT64, 4]\n",
      "  %onnx::MatMul_197[FLOAT, 512x512]\n",
      ") {\n",
      "  %onnx::Reshape_196 = Identity(%onnx::Reshape_176)\n",
      "  %onnx::Reshape_191 = Identity(%onnx::Reshape_171)\n",
      "  %onnx::Reshape_186 = Identity(%onnx::Reshape_176)\n",
      "  %onnx::Reshape_181 = Identity(%onnx::Reshape_171)\n",
      "  %tensor = MatMul(%onnx::MatMul_0, %onnx::MatMul_164)\n",
      "  %tensor.3 = MatMul(%onnx::MatMul_0, %onnx::MatMul_165)\n",
      "  %tensor.7 = MatMul(%onnx::MatMul_0, %onnx::MatMul_166)\n",
      "  %tensor.11 = Reshape(%tensor, %onnx::Reshape_171)\n",
      "  %tensor.15 = Transpose[perm = [0, 2, 1, 3]](%tensor.11)\n",
      "  %x = Reshape(%tensor.15, %onnx::Reshape_176)\n",
      "  %tensor.19 = Reshape(%tensor.3, %onnx::Reshape_181)\n",
      "  %tensor.23 = Transpose[perm = [0, 2, 1, 3]](%tensor.19)\n",
      "  %x.3 = Reshape(%tensor.23, %onnx::Reshape_186)\n",
      "  %tensor.27 = Reshape(%tensor.7, %onnx::Reshape_191)\n",
      "  %tensor.31 = Transpose[perm = [0, 2, 1, 3]](%tensor.27)\n",
      "  %x.7 = Reshape(%tensor.31, %onnx::Reshape_196)\n",
      "  %onnx::Add_91 = Einsum[equation = '...ij,...kj->...ik'](%x, %x.3)\n",
      "  %onnx::Add_92 = Constant[value = <Tensor>]()\n",
      "  %onnx::Div_93 = Add(%onnx::Add_92, %onnx::Add_91)\n",
      "  %onnx::Pow_94 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Pow_95 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Div_96 = Pow(%onnx::Pow_94, %onnx::Pow_95)\n",
      "  %onnx::Softmax_97 = Div(%onnx::Div_93, %onnx::Div_96)\n",
      "  %x.11 = Softmax[axis = -1](%onnx::Softmax_97)\n",
      "  %tensor.35 = Einsum[equation = 'abde,abcd->abce'](%x.7, %x.11)\n",
      "  %onnx::Gather_100 = Shape(%tensor.35)\n",
      "  %onnx::Gather_101 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Div_102 = Gather[axis = 0](%onnx::Gather_100, %onnx::Gather_101)\n",
      "  %onnx::Gather_103 = Shape(%tensor.35)\n",
      "  %onnx::Gather_104 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Div_105 = Gather[axis = 0](%onnx::Gather_103, %onnx::Gather_104)\n",
      "  %onnx::Gather_106 = Shape(%tensor.35)\n",
      "  %onnx::Gather_107 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Div_108 = Gather[axis = 0](%onnx::Gather_106, %onnx::Gather_107)\n",
      "  %onnx::Gather_109 = Shape(%tensor.35)\n",
      "  %onnx::Gather_110 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Div_111 = Gather[axis = 0](%onnx::Gather_109, %onnx::Gather_110)\n",
      "  %onnx::Div_112 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Cast_113 = Div(%onnx::Div_102, %onnx::Div_112)\n",
      "  %onnx::Cast_114 = Cast[to = 7](%onnx::Cast_113)\n",
      "  %onnx::Unsqueeze_115 = Cast[to = 7](%onnx::Cast_114)\n",
      "  %onnx::Div_116 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Cast_117 = Div(%onnx::Div_105, %onnx::Div_116)\n",
      "  %onnx::Cast_118 = Cast[to = 7](%onnx::Cast_117)\n",
      "  %onnx::Mul_119 = Cast[to = 7](%onnx::Cast_118)\n",
      "  %onnx::Div_120 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Cast_121 = Div(%onnx::Div_108, %onnx::Div_120)\n",
      "  %onnx::Cast_122 = Cast[to = 7](%onnx::Cast_121)\n",
      "  %onnx::Unsqueeze_123 = Cast[to = 7](%onnx::Cast_122)\n",
      "  %onnx::Div_124 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Cast_125 = Div(%onnx::Div_111, %onnx::Div_124)\n",
      "  %onnx::Cast_126 = Cast[to = 7](%onnx::Cast_125)\n",
      "  %inferred_length.11 = Cast[to = 7](%onnx::Cast_126)\n",
      "  %onnx::Mul_128 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Mul_129 = Mul(%onnx::Mul_119, %onnx::Mul_128)\n",
      "  %onnx::Unsqueeze_130 = Mul(%onnx::Mul_129, %inferred_length.11)\n",
      "  %onnx::Unsqueeze_131 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_132 = Unsqueeze(%onnx::Unsqueeze_115, %onnx::Unsqueeze_131)\n",
      "  %onnx::Unsqueeze_133 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_134 = Unsqueeze(%onnx::Mul_119, %onnx::Unsqueeze_133)\n",
      "  %onnx::Unsqueeze_135 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_136 = Unsqueeze(%onnx::Unsqueeze_123, %onnx::Unsqueeze_135)\n",
      "  %onnx::Unsqueeze_137 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_138 = Unsqueeze(%inferred_length.11, %onnx::Unsqueeze_137)\n",
      "  %onnx::Reshape_139 = Concat[axis = 0](%onnx::Concat_132, %onnx::Concat_134, %onnx::Concat_136, %onnx::Concat_138)\n",
      "  %tensor.39 = Reshape(%tensor.35, %onnx::Reshape_139)\n",
      "  %tensor.43 = Transpose[perm = [0, 2, 1, 3]](%tensor.39)\n",
      "  %onnx::Unsqueeze_142 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_143 = Unsqueeze(%onnx::Unsqueeze_115, %onnx::Unsqueeze_142)\n",
      "  %onnx::Unsqueeze_144 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_145 = Unsqueeze(%onnx::Unsqueeze_123, %onnx::Unsqueeze_144)\n",
      "  %onnx::Unsqueeze_146 = Constant[value = <Tensor>]()\n",
      "  %onnx::Concat_147 = Unsqueeze(%onnx::Unsqueeze_130, %onnx::Unsqueeze_146)\n",
      "  %onnx::Reshape_148 = Concat[axis = 0](%onnx::Concat_143, %onnx::Concat_145, %onnx::Concat_147)\n",
      "  %onnx::MatMul_149 = Reshape(%tensor.43, %onnx::Reshape_148)\n",
      "  %onnx::Add_151 = MatMul(%onnx::MatMul_149, %onnx::MatMul_197)\n",
      "  %input = Add(%onnx::MatMul_0, %onnx::Add_151)\n",
      "  %onnx::Sub_153 = ReduceMean[axes = [-1]](%input)\n",
      "  %onnx::Pow_154 = Sub(%input, %onnx::Sub_153)\n",
      "  %onnx::Pow_155 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::ReduceMean_156 = Pow(%onnx::Pow_154, %onnx::Pow_155)\n",
      "  %onnx::Add_157 = ReduceMean[axes = [-1]](%onnx::ReduceMean_156)\n",
      "  %onnx::Add_158 = Constant[value = <Scalar Tensor []>]()\n",
      "  %onnx::Sqrt_159 = Add(%onnx::Add_157, %onnx::Add_158)\n",
      "  %onnx::Div_160 = Sqrt(%onnx::Sqrt_159)\n",
      "  %onnx::Mul_161 = Div(%onnx::Pow_154, %onnx::Div_160)\n",
      "  %onnx::Add_162 = Mul(%onnx::Mul_161, %norm.weight)\n",
      "  %163 = Add(%onnx::Add_162, %norm.bias)\n",
      "  return %163\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(onnx.helper.printable_graph(onnx_model.graph))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "701a749a-8b2e-4e6c-9e83-11dfffc4555a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime as ort\n",
    "ort_session = ort.InferenceSession(\"mha.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f3843b1c-da7a-4c55-be7e-3a6f1782aad6",
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = {ort_session.get_inputs()[0].name: emb.detach().numpy()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "38af0c5b-1c11-4a1c-88ef-cc7b6408d68c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 22.6 s, sys: 4.31 s, total: 26.9 s\n",
      "Wall time: 5.41 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "out = ort_session.run(None, inp) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "65651a88-ff0f-4e2a-8bb9-7e8c783f0939",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-08-06 17:13:57.261734 [W:onnxruntime:, inference_session.cc:1488 Initialize] Serializing optimized model with Graph Optimization level greater than ORT_ENABLE_EXTENDED and the NchwcTransformer enabled. The generated model may contain hardware specific optimizations, and should only be used in the same environment the model was optimized in.\n"
     ]
    }
   ],
   "source": [
    "sess_options = ort.SessionOptions()\n",
    "sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL\n",
    "sess_options.optimized_model_filepath = \"mha_optim.onnx\"\n",
    "session = ort.InferenceSession(\"mha.onnx\", sess_options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f17437a9-de92-48b6-a967-16b1d073890e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 22.9 s, sys: 4.62 s, total: 27.6 s\n",
      "Wall time: 6.48 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "out = session.run(None, inp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6f23a57-511a-44c6-82a7-efdf9f098abe",
   "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}