vision_transformers.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b2434719",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from einops.layers.torch import Rearrange\n",
    "from einops import repeat, rearrange"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ce52931",
   "metadata": {},
   "source": [
    "## Resources \n",
    "- [Vision Transformers](https://arxiv.org/pdf/2010.11929.pdf)\n",
    "- [Attention by lucidrains](https://github.com/lucidrains/vit-pytorch/blob/4b8f5bc90002a5506d765c811b554760d8dd6ee7/vit_pytorch/vit.py#L35)\n",
    "- [Transformers by lucidrains](https://github.com/lucidrains/vit-pytorch/blob/4b8f5bc90002a5506d765c811b554760d8dd6ee7/vit_pytorch/vit.py#L67)\n",
    "- [Attention is all you need](https://arxiv.org/pdf/1706.03762.pdf)\n",
    "\n",
    "\n",
    "Lets understand `Vision Transformer` in 6 simple steps "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab70135d",
   "metadata": {},
   "source": [
    "## step-1\n",
    "The following are the inputs required by the vision transformer. \n",
    "- input image size. \n",
    "- patch_size \n",
    "\n",
    "from this we can calculate the number of patches and patch_dimension in the following way\n",
    "\n",
    "> If the image is of size (H, W) and our patch size (PxP). we reshape the image to $I^{N, P^2xC}$ where C is the number of channels and N= HxW/$P^2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c7a1a455",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = torch.randn((224*224*3)).reshape((224, 224, 3))\n",
    "image_shape = image.shape\n",
    "patch_size = (14, 14)\n",
    "\n",
    "image_height, image_width, channels = image_shape\n",
    "patch_height, patch_width = patch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "117af387",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total_patches: 256\n",
      "patch_dim: 588\n"
     ]
    }
   ],
   "source": [
    "num_patches = (image_height // patch_height) * (image_width // patch_width)\n",
    "patch_dim = channels * patch_height * patch_width\n",
    "\n",
    "print(f\"total_patches: {num_patches}\")\n",
    "print(f\"patch_dim: {patch_dim}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3870eee7",
   "metadata": {},
   "source": [
    "> Reshape the input image to [batch_size, total_patches, patch_dim]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "35c05de3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 256, 588])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mr = Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width)\n",
    "out = mr(image.permute((2, 0, 1)).unsqueeze(0))\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "365e0719",
   "metadata": {},
   "source": [
    "## step-2\n",
    "> The Transformer uses constant latent vector size D through all of its layers, so we flatten the patches and map to D dimensions with a trainable linear projection. We refer to the output of this projection as the `patch embeddings`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0dfb5b1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 256, 128])\n"
     ]
    }
   ],
   "source": [
    "embed_dim = 128\n",
    "embed = torch.nn.Linear(588, 128)\n",
    "with torch.no_grad():\n",
    "    embed_out = embed(out)\n",
    "print(embed_out.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6781bcad",
   "metadata": {},
   "source": [
    "## step-3\n",
    "> Similar to BERT’s [class] token, we prepend a learnable embedding to the sequence of embedded patches ($z^0_{0}$ = xclass), whose state at the output of the Transformer encoder ($z^0_{L}$) serves as the image representation y (Eq. 4). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a870646c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 128])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cls_token = torch.nn.Parameter(torch.randn(1, 1, embed_dim))\n",
    "cls_token.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41fc4a21",
   "metadata": {},
   "source": [
    "we have to repeat this `cls_token` batch_size times and concat with our `step-2` output\n",
    "> we can use one of `torch.repeat` or `eniops repeat` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "438fff99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 316 µs, sys: 107 µs, total: 423 µs\n",
      "Wall time: 284 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 128])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "b, n, _ = embed_out.shape\n",
    "cls_tokens = cls_token.repeat(b, 1, 1)\n",
    "cls_tokens.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8f6b9126",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 779 µs, sys: 577 µs, total: 1.36 ms\n",
      "Wall time: 901 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 128])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "b, n, _ = embed_out.shape\n",
    "cls_tokens = repeat(cls_token, '1 1 d -> b 1 d', b = b) #we repeat N of batch times \n",
    "cls_tokens.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f7719b4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.cat((cls_tokens, embed_out), dim=1)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acca8f81",
   "metadata": {},
   "source": [
    "## step4\n",
    "> Position embeddings are added to the patch embeddings to retain positional information. We use standard learnable 1D position embeddings, The resulting sequence of embedding vectors serves as input to the encoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "389b27a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pos_embedding = torch.nn.Parameter(torch.randn(1, num_patches + 1, embed_dim))\n",
    "pos_embedding.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6700b9af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x += pos_embedding[:, :(n + 1)]\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59bfd1a1",
   "metadata": {},
   "source": [
    "> Add dropout if required. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d97469e8",
   "metadata": {},
   "source": [
    "## step-5\n",
    "How a transformer works?\n",
    "- [transformer is self-attention, multi-head] x blocks\n",
    "- each transformer block is shown below\n",
    "- it has two norm layer, one attention layer and one feed-forward layer. residual is applied after attention and after `ffn`\n",
    "\n",
    "To be more specific\n",
    "> 5.1) `y1 = Norm(x)`\n",
    "\n",
    "> 5.2) `z1 = Attention(y1)`\n",
    "\n",
    "> 5.3) `z1=+x (Residual layer)`\n",
    "\n",
    "> 5.4) `y = Norm(z1)`\n",
    "\n",
    "> 5.5) `z = FF(y)`\n",
    "\n",
    "> 5.6) `x= z+z1 (residual layer)`\n",
    "\n",
    "> 5.7) combine everything into one transformer block output is x which is sent to next res block\n",
    "\n",
    "\n",
    "<img src=\"../images/transformer_block.png\" alt=\"alt text\" width=\"125\" align=\"left\"/>\n",
    "\n",
    "we will apply each step here. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7687a7a4",
   "metadata": {},
   "source": [
    "### step 5.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "39118006",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norm1 = torch.nn.LayerNorm(embed_dim)\n",
    "\n",
    "with torch.no_grad():\n",
    "    y1 = norm1(x)\n",
    "y1.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1174c5c",
   "metadata": {},
   "source": [
    "### 5.2 Attention\n",
    "This is the most important aspect of Transformer. It involves two parts\n",
    "- self-attention\n",
    "- multi-head\n",
    "\n",
    "Attention block has the following attributes\n",
    "- dim \n",
    "- heads: total number of heads\n",
    "- dim_head: dimension of each head\n",
    "- dropout: optional after `qk`\n",
    "\n",
    "\n",
    "#### self-attention\n",
    "- we take the `norm vector` obtained in 5.1 and linearly project this tensor into 3 tensor `q`, `k`, `v` using a feed forward neural network. \n",
    "- then self-attention is just a dot product of q and k. scaled by $\\sqrt(d_{k}$. softmax is applied on this output. The output again is a dot product with v. The [vaswani et al](https://arxiv.org/pdf/1706.03762.pdf) has insights (section 3.1) on \n",
    "- why dot product is choosen over additive?\n",
    "- why we scale the outputs of qk? \n",
    "\n",
    "$$\n",
    "Attention(Q, K, V ) = softmax(QK^{T}/\\sqrt{d_{k}})V\n",
    "$$\n",
    "\n",
    "<img src=\"../images/attn_qkv.png\" alt=\"alt text\" width=\"200\" align=\"left\"/>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9eb916df",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionHead(torch.nn.Module):\n",
    "    def __init__(self, dim, dim_head, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.dim = dim \n",
    "        self.dim_head = dim_head\n",
    "        self.scale = dim_head**(-0.5)\n",
    "        self.attend = torch.nn.Softmax(dim = -1)\n",
    "        self.to_qkv = torch.nn.Linear(self.dim, self.dim_head*3, bias = False)\n",
    "        self.drop = torch.nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        qkv = self.to_qkv(x).chunk(3, dim=-1)\n",
    "        q, k, v = qkv\n",
    "        qk = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n",
    "        qk_soft = self.attend(qk)\n",
    "        qk_soft = self.drop(qk_soft)\n",
    "        out = torch.matmul(qk_soft, v)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9f9de0c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 64])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ah = AttentionHead(128, 64)\n",
    "ah(y1).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5752e20d",
   "metadata": {},
   "source": [
    "<img src=\"../images/attn_multihead.png\" alt=\"alt text\" width=\"200\" align=\"left\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "25ef6591",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionAllHead(torch.nn.Module):\n",
    "    def __init__(self, dim, heads, dim_head, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.heads = heads\n",
    "        project_out = not (heads == 1 and dim_head == dim)\n",
    "        for i in range(self.heads):\n",
    "            setattr(self, f\"head_{i}\", AttentionHead(dim, dim_head, dropout))\n",
    "        #self.attnhead = [AttentionHead(dim, dim_head, dropout) for i in range(self.heads)]\n",
    "        inner_dim = self.heads * dim_head\n",
    "        self.to_out = torch.nn.Sequential(\n",
    "            torch.nn.Linear(inner_dim, dim),\n",
    "            torch.nn.Dropout(dropout)\n",
    "        ) if project_out else nn.Identity()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        dd = []\n",
    "        for i in range(self.heads):\n",
    "            out = getattr(self, f\"head_{i}\")(x)\n",
    "            dd.append(out.unsqueeze(1))\n",
    "        #dd = [t(x).unsqueeze(1) for t in self.attnhead]\n",
    "        dd = torch.cat(dd, dim=1)\n",
    "        out = rearrange(dd, 'b h n d -> b n (h d)')\n",
    "        return self.to_out(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a98fbf5",
   "metadata": {},
   "source": [
    "> I have kept the implementation little simple. this [repo](https://github.com/lucidrains/vit-pytorch/blob/4b8f5bc90002a5506d765c811b554760d8dd6ee7/vit_pytorch/vit.py#L47) has combined both our functions into one called `Attention` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b761f83d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dim = 128\n",
    "heads = 8\n",
    "dim_head = 64\n",
    "alh = AttentionAllHead(128, 8, 64, 0.0)\n",
    "z1 = alh(y1)\n",
    "z1.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99196c9d",
   "metadata": {},
   "source": [
    "### 5.3 residual "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "02d61622",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z1+=x \n",
    "z1.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6763f74c",
   "metadata": {},
   "source": [
    "### 5.4 Norm2 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e77d0631",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norm2 = torch.nn.LayerNorm(embed_dim)\n",
    "\n",
    "with torch.no_grad():\n",
    "    y = norm2(z1)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "955feb5e",
   "metadata": {},
   "source": [
    "### 5.5 Feadforward neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7197d3df",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(torch.nn.Module):\n",
    "    def __init__(self, dim, hidden_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.net = torch.nn.Sequential(\n",
    "            torch.nn.Linear(dim, hidden_dim),\n",
    "            torch.nn.GELU(),\n",
    "            torch.nn.Dropout(dropout),\n",
    "            torch.nn.Linear(hidden_dim, dim),\n",
    "            torch.nn.Dropout(dropout)\n",
    "        )\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "44f3efe1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff = FeedForward(128, hidden_dim=2048, dropout=0.1)\n",
    "with torch.no_grad():\n",
    "    z = ff(y)\n",
    "z.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fa5594c",
   "metadata": {},
   "source": [
    "### 5.6 Another residual layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "21f267fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 257, 128])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = z+z1\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d94a9dd",
   "metadata": {},
   "source": [
    "> Combine everything"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05622d2f",
   "metadata": {},
   "source": [
    "## step-6 Transformer block "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e1f30946",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(torch.nn.Module):\n",
    "    def __init__(self, dim, heads, dim_head, mlp_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.norm1 = torch.nn.LayerNorm(dim)\n",
    "        self.norm2 = torch.nn.LayerNorm(dim)\n",
    "        self.attn = AttentionAllHead(dim, heads, dim_head, dropout)\n",
    "        self.ff = FeedForward(dim, hidden_dim=mlp_dim, dropout=dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        y1 = self.norm1(x)\n",
    "        z1 = self.attn(y1)\n",
    "        z1+=x\n",
    "        y = self.norm2(z1)\n",
    "        z = self.ff(y)\n",
    "        x = z+z1\n",
    "        return x "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "18a49a4e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 256, 128])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tb = TransformerBlock(dim=128, heads=8, dim_head=64, mlp_dim=2048)\n",
    "tb(embed_out).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffbf96b0",
   "metadata": {},
   "source": [
    "> The vision transformer uses 6 transformer blocks. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c782e72",
   "metadata": {},
   "source": [
    "we can use `vit_pytorch` to load a transformer network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "be9197fb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1000])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from vit_pytorch import ViT\n",
    "\n",
    "v = ViT(\n",
    "    image_size = 256,\n",
    "    patch_size = 32,\n",
    "    num_classes = 1000,\n",
    "    dim = 1024,\n",
    "    depth = 6,\n",
    "    heads = 16,\n",
    "    mlp_dim = 2048,\n",
    "    dropout = 0.1,\n",
    "    emb_dropout = 0.1\n",
    ")\n",
    "\n",
    "img = torch.randn(1, 3, 256, 256)\n",
    "\n",
    "preds = v(img) # (1, 1000)\n",
    "preds.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2de06bd6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "54622184"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum([params.numel() for name, params in v.named_parameters()])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9550a59c",
   "metadata": {},
   "source": [
    "> ViT has 54 million params :( "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (cv)",
   "language": "python",
   "name": "cv"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}