ViT/building_ViT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "181123bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from types import SimpleNamespace\n",
    "from einops import einsum, rearrange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4752784b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "namespace(embed_dim=128,\n",
       "          num_heads=4,\n",
       "          depth=6,\n",
       "          pool='mean',\n",
       "          img_size=224,\n",
       "          num_channels=3,\n",
       "          patch_size=16,\n",
       "          attention_dropout=0.0,\n",
       "          residual_dropout=0.0,\n",
       "          mlp_ratio=4,\n",
       "          mlp_dropout=0.0,\n",
       "          pos_dropout=0.0,\n",
       "          num_classes=1000)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config_vit_pico = SimpleNamespace(\n",
    "    embed_dim = 128,\n",
    "    num_heads = 4,\n",
    "    depth = 6,\n",
    "    pool = 'mean',\n",
    "    img_size = 224,\n",
    "    num_channels = 3,\n",
    "    patch_size = 16,\n",
    "    attention_dropout = 0.,\n",
    "    residual_dropout = 0.,\n",
    "    mlp_ratio = 4,\n",
    "    mlp_dropout = 0.,\n",
    "    pos_dropout = 0.,\n",
    "    num_classes = 1000\n",
    ")\n",
    "\n",
    "config_vit_pico"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e67c59ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def params(m):\n",
    "    return sum([p.numel() for p in m.parameters() if p.requires_grad])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "97a4b463",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self,config):\n",
    "        super().__init__()\n",
    "        self.embed_dim = config.embed_dim\n",
    "        self.n_heads = config.num_heads\n",
    "        assert self.embed_dim % self.n_heads == 0, 'embedding dimension by be divisible by number of heads'\n",
    "        self.head_size = self.embed_dim // self.n_heads\n",
    "        self.seq_len = config.embed_dim\n",
    "        \n",
    "        self.qkv = nn.Linear(self.embed_dim, self.head_size * self.n_heads * 3,bias=False)\n",
    "        self.scale = self.head_size ** -0.5\n",
    "        \n",
    "        self.attention_dropout = nn.Dropout(config.attention_dropout)\n",
    "        \n",
    "        self.proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)\n",
    "        self.residual_dropout = nn.Dropout(config.residual_dropout)\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        b,t,c = x.shape\n",
    "        # q,k,v shape individually: batch_size x seq_len x embed_dim\n",
    "        # we know that qk_t = q x k_t, where q=bxtxhead_dim, k_t=bxhead_timxt\n",
    "        q,k,v = self.qkv(x).chunk(3,dim=-1)\n",
    "        q = rearrange(q,'b t (h n) -> b n t h',n=self.n_heads) # h = head_size\n",
    "        k = rearrange(k,'b t (h n) -> b n t h',n=self.n_heads)\n",
    "        v = rearrange(v,'b t (h n) -> b n t h',n=self.n_heads)\n",
    "        \n",
    "        # qk_t = einsum(q,k,'b n t1 h, b n t2 h -> b n t1 t2') * self.scale\n",
    "        qk_t = (q@k.transpose(-2,-1)) * self.scale\n",
    "        \n",
    "        weights = self.attention_dropout(qk_t)\n",
    "        \n",
    "        attention = weights @ v # batch x n_heads x seq_len x head_size\n",
    "        attention = rearrange(attention,'b n t h -> b t (n h)') # batch x n_heads x seq_len x embed_dim\n",
    "        \n",
    "        out = self.proj(attention)\n",
    "        out = self.residual_dropout(out)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "92aceb79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 128, 128])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.rand(16,config_vit_pico.embed_dim,config_vit_pico.embed_dim)\n",
    "a = MultiHeadAttention(config_vit_pico)\n",
    "a(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9db10faa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self,config):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(config.embed_dim)\n",
    "        self.attn = MultiHeadAttention(config)\n",
    "        self.ln2 = nn.LayerNorm(config.embed_dim)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(config.embed_dim,config.embed_dim*config.mlp_ratio),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(config.embed_dim*config.mlp_ratio,config.embed_dim),\n",
    "            nn.Dropout(config.mlp_dropout)\n",
    "        )\n",
    "        \n",
    "    def forward(self,x):\n",
    "        x = x+self.attn(self.ln1(x))\n",
    "        x = x+self.mlp(self.ln2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a0785db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ViT(nn.Module):\n",
    "    def __init__(self,config):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        config.num_patches = (config.img_size // config.patch_size) ** 2\n",
    "        config.patch_dim = config.num_channels * config.patch_size ** 2\n",
    "        \n",
    "        self.config = config\n",
    "        \n",
    "        self.patch_embedding = nn.Sequential(\n",
    "            nn.LayerNorm(self.config.patch_dim),\n",
    "            nn.Linear(self.config.patch_dim, self.config.embed_dim, bias=False),\n",
    "            nn.LayerNorm(self.config.embed_dim)\n",
    "        )\n",
    "        self.pos_embed = nn.Parameter(torch.randn(1,self.config.num_patches,self.config.embed_dim),requires_grad=True)\n",
    "        self.pos_dropout = nn.Dropout(self.config.pos_dropout)\n",
    "        \n",
    "        self.transformer_blocks = nn.ModuleList([TransformerBlock(config) for _ in range(self.config.depth)])\n",
    "        \n",
    "        self.head = nn.Linear(self.config.embed_dim,self.config.num_classes)\n",
    "        \n",
    "    def forward(self,x):\n",
    "        \n",
    "        x = rearrange(x,'b c (h p1) (w p2) -> b (h w) (p1 p2 c)',\n",
    "                      p1=self.config.patch_size,\n",
    "                      p2=self.config.patch_size\n",
    "                     )\n",
    "        x = self.patch_embedding(x)\n",
    "        x += self.pos_embed\n",
    "        x = self.pos_dropout(x)\n",
    "        \n",
    "        for block in self.transformer_blocks:\n",
    "            x = block(x)\n",
    "            \n",
    "        x = x.mean(dim=1)\n",
    "        x = self.head(x)\n",
    "        \n",
    "        return x\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0db2a1bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 1000]),\n",
       " 1440744,\n",
       " namespace(embed_dim=128,\n",
       "           num_heads=4,\n",
       "           depth=6,\n",
       "           pool='mean',\n",
       "           img_size=224,\n",
       "           num_channels=3,\n",
       "           patch_size=16,\n",
       "           attention_dropout=0.0,\n",
       "           residual_dropout=0.0,\n",
       "           mlp_ratio=4,\n",
       "           mlp_dropout=0.0,\n",
       "           pos_dropout=0.0,\n",
       "           num_classes=1000,\n",
       "           num_patches=196,\n",
       "           patch_dim=768))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = ViT(config_vit_pico)\n",
    "x = torch.rand(1,3,224,224)\n",
    "model(x).shape, params(model), model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2aeeecc5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ViT(\n",
       "  (patch_embedding): Sequential(\n",
       "    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (1): Linear(in_features=768, out_features=128, bias=False)\n",
       "    (2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (pos_dropout): Dropout(p=0.0, inplace=False)\n",
       "  (transformer_blocks): ModuleList(\n",
       "    (0-5): 6 x TransformerBlock(\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): MultiHeadAttention(\n",
       "        (qkv): Linear(in_features=128, out_features=384, bias=False)\n",
       "        (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=False)\n",
       "        (residual_dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): Sequential(\n",
       "        (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "        (1): GELU(approximate='none')\n",
       "        (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "        (3): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (head): Linear(in_features=128, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7f7bd477",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 1000]), 86530024)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config_vit_base = SimpleNamespace(\n",
    "    embed_dim = 768,\n",
    "    num_heads = 12,\n",
    "    depth = 12,\n",
    "    img_size = 224,\n",
    "    num_channels = 3,\n",
    "    patch_size = 16,\n",
    "    attention_dropout = 0.,\n",
    "    residual_dropout = 0.,\n",
    "    mlp_ratio = 4,\n",
    "    mlp_dropout = 0.,\n",
    "    pos_dropout = 0.,\n",
    "    num_classes = 1000\n",
    ")\n",
    "\n",
    "base = ViT(config_vit_base)\n",
    "x = torch.rand(1,3,224,224)\n",
    "base(x).shape, params(base)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7b0e8d41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ViT(\n",
       "  (patch_embedding): Sequential(\n",
       "    (0): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (1): Linear(in_features=768, out_features=768, bias=False)\n",
       "    (2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (pos_dropout): Dropout(p=0.0, inplace=False)\n",
       "  (transformer_blocks): ModuleList(\n",
       "    (0-11): 12 x TransformerBlock(\n",
       "      (ln1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): MultiHeadAttention(\n",
       "        (qkv): Linear(in_features=768, out_features=2304, bias=False)\n",
       "        (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "        (residual_dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (ln2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): Sequential(\n",
       "        (0): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (1): GELU(approximate='none')\n",
       "        (2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (3): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (head): Linear(in_features=768, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0b447074",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-3.2585e-01, -2.4704e+00, -1.4879e-01,  ..., -6.1265e-01,\n",
       "           1.1498e+00, -6.2170e-02],\n",
       "         [-1.3894e+00,  4.6732e-01, -1.8328e-03,  ...,  1.2957e+00,\n",
       "           1.9182e-01, -1.8384e-01],\n",
       "         [ 8.2649e-01, -2.5241e-01,  4.8311e-01,  ..., -1.2235e+00,\n",
       "           4.6463e-01,  1.5737e-01],\n",
       "         ...,\n",
       "         [-1.5943e-01,  1.4821e+00,  8.3911e-01,  ..., -6.4295e-01,\n",
       "          -1.0655e+00,  8.7774e-01],\n",
       "         [ 2.2364e+00, -8.3828e-01, -2.2446e+00,  ...,  6.1022e-02,\n",
       "          -1.2066e+00,  1.0642e+00],\n",
       "         [ 5.2472e-01,  5.3328e-02, -6.7582e-01,  ...,  6.5444e-01,\n",
       "           2.0783e-03, -1.1407e+00]]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.randn(1, 128, 128)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d12ebf",
   "metadata": {},
   "source": [
    "# fix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "28a135b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand(1,3,128,128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "69d2a967",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(16, 64, 768)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "patch_size = 16\n",
    "num_patches = ((128//16)**2)\n",
    "patch_embed = 3 * 16**2\n",
    "patch_size,num_patches, patch_embed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "402a2013",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 768, 8, 8])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "patch_conv = nn.Conv2d(3,patch_embed,16,16)\n",
    "x2 = patch_conv(x)\n",
    "x2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "72dc974e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64, 768])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x2 = rearrange(x2,'b p h w -> b (h w) p')\n",
    "x2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "2fbdcb6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64, 64])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "patchnorm = nn.Sequential(\n",
    "    nn.LayerNorm(patch_embed),\n",
    "    nn.Linear(patch_embed,64),\n",
    "    nn.LayerNorm(64)\n",
    ")\n",
    "x3 = patchnorm(x2)\n",
    "x3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "fafdd980",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64, 64])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pos = torch.randn(1,num_patches,64)\n",
    "x3 += pos\n",
    "x3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96ba4b87",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}