25Nov-Exper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0faa17de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a99c4302",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: einops in /opt/conda/lib/python3.8/site-packages (0.3.2)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cfef700e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2e000007",
   "metadata": {},
   "outputs": [],
   "source": [
    "from einops import rearrange\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2a025aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PatchEmbed(nn.Module):\n",
    "    \n",
    "    \"\"\"\n",
    "    To split images into patches of given size, and convert to embeddings\n",
    "    \n",
    "    Parameters:\n",
    "    --------\n",
    "    img_size: int   (Size of the image)\n",
    "    \n",
    "    patch_size: int  (Size of Patch)\n",
    "    \n",
    "    in_chans: int (No of input channels)\n",
    "    \n",
    "    embedding_dims: int (The embeddings dimension)\n",
    "    \n",
    "    \n",
    "    Attributes:\n",
    "    -------\n",
    "    n_patches : int   (No of patches inside the image)\n",
    "    \n",
    "    proj: nn.Conv2d\n",
    "        Convolution Layer that do both the splitting into patches, and their embeddings\n",
    "        \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,img_size,patch_size,in_chans=3,embed_dims=768):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.img_size=img_size\n",
    "        self.patch_size=patch_size\n",
    "        self.n_patches=(img_size//patch_size)**2\n",
    "        \n",
    "        \n",
    "        self.proj=nn.Conv2d(in_chans,embed_dims,\n",
    "                           kernel_size=patch_size,\n",
    "                           stride=patch_size\n",
    "                           )\n",
    "        \n",
    "    def forward(self,x):\n",
    "        \n",
    "        \"\"\"\n",
    "        parameters:\n",
    "        ---------\n",
    "        x: torch.Tensor   Shape '(n_samples,in_chans, img_size,img_size)'\n",
    "        \n",
    "        Returns\n",
    "        --------\n",
    "        torch.Tensor  Shape '(n_samples,n_patches, embed_dims)'\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        x=self.proj(\n",
    "                    x\n",
    "                    )  # (n_samples,embed_dims,n_patches ** 0.5, n_patches ** 0.5)\n",
    "        x=x.flatten(2)\n",
    "        x=x.transpose(1,2) # (n_samples,n_patches, embed_dims)\n",
    "        \n",
    "        \n",
    "        return x\n",
    "    \n",
    "    \n",
    "        \n",
    "        \n",
    "        \n",
    "     "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f83e947",
   "metadata": {},
   "source": [
    "# Experimental"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c7978a5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 5, 5]), torch.Size([3, 7, 4]))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#to del\n",
    "x=torch.ones(3,5,5)\n",
    "x2=torch.ones(3,7,4)\n",
    "x.shape,x2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2bb20bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_conv=nn.Conv2d(3,768,2,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f3857581",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 768, 2, 2])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out=test_conv(x.unsqueeze(0))\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "dd386f9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 768, 3, 2])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out2=test_conv(x2.unsqueeze(0))\n",
    "out2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "628ddae4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 768, 6])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out2=out2.flatten(2)\n",
    "out2.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dce3dea",
   "metadata": {},
   "source": [
    "* **So we can also design transformers to work with rectangle images**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5db2030e",
   "metadata": {},
   "source": [
    "# Implementing Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d179ca28",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    \"\"\"\n",
    "    Attention Mechanism\n",
    "    \n",
    "    Parameters:\n",
    "    -------\n",
    "    dim: int  Input and output ims of per patch(or token) features\n",
    "    \n",
    "    n_heads: int   No of attention heads\n",
    "    \n",
    "    qkv_bias: True  If True, then we include bias to the query, key, and value projections\n",
    "    \n",
    "    attn_dp: float  Dropout prob applied to the k,q,and v tensors\n",
    "    \n",
    "    fin_proj_dp: float Dropout prob applied to the output tensor\n",
    "    \n",
    "    Attributes:\n",
    "    --------\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,dim,n_heads=12,qkv_bias=True, attn_dp=0.0, fin_proj_dp=0.0):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.dim=dim\n",
    "        self.n_heads=n_heads\n",
    "        self.head_dim= dim//n_heads\n",
    "        self.scale=self.head_dim ** -0.5\n",
    "        \n",
    "        self.qkv=nn.Linear(dim,dim*3,bias=qkv_bias)\n",
    "        self.attn_dp=nn.Dropout(attn_dp)\n",
    "        self.fin_proj=nn.Linear(dim,dim)\n",
    "        self.fin_proj_dp=nn.Dropout(fin_proj_dp)\n",
    "        \n",
    "    def forward(self,x):\n",
    "        \n",
    "        \"\"\"\n",
    "        Run Forward Pass\n",
    "        \"\"\"\n",
    "        \n",
    "        n_samples,n_tokens,dim=x.shape\n",
    "        qkv=self.qkv(x) # (n_samples,n_patches+1,dim*3)\n",
    "        qkv=qkv.reshape(\n",
    "                        n_samples,n_tokens,3,self.n_heads,self.head_dim\n",
    "        )\n",
    "        qkv=qkv.permute(2,0,3,1,4)  # (3,n_samples,n_heads,n_patches+1,head_dim)\n",
    "        k_t=k.transpose(-2,-1) # (n_samples,n_heads,head_dim,n_patches+1)\n",
    "        dp=(q @ k_t) * self.scale   # (n_samples,n_heads, n_patches+1, n_patches+1)\n",
    "        attn=dp.softmax(dim=-1)\n",
    "        attn=self.attn_dp(attn)\n",
    "        \n",
    "        weighted_avg=attn @ v   # (n_samples,n_heads, n_patches+1, heads_dim)\n",
    "        weighted_avg=weighted_avg.transpose(1,2)  # (n_samples,n_patches+1,n_heads,heads_dim)\n",
    "        \n",
    "        weighted_avg=weighted_avg.flatten(2)  # (n_samples,n_patches+1,dim)\n",
    "        x=self.fin_proj(weighted_avg)\n",
    "        x=self.fin_proj_dp(x)\n",
    "        \n",
    "        # So we returned the same shape as input was of.\n",
    "        return x\n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533f1fc9",
   "metadata": {},
   "source": [
    "# Writing MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "af5b3d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    \n",
    "    def __init__(self,in_features,hidden_features,out_features,p=0.):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.fc1=nn.Linear(in_features,hidden_features)\n",
    "        self.act=nn.GELU()\n",
    "        self.fc2=nn.Linear(hidden_features,out_features)\n",
    "        self.drop=nn.Dropout(p)\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"Run forward pass.\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Shape `(n_samples, n_patches + 1, in_features)`.\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Shape `(n_samples, n_patches +1, out_features)`\n",
    "        \"\"\"\n",
    "        x = self.fc1(\n",
    "                x\n",
    "        ) # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.act(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.drop(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.fc2(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.drop(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "79ca4dc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, dim, n_heads, mlp_ratio=4.0, qkv_bias=True, p=0., attn_p=0.):\n",
    "        super().__init__()\n",
    "        self.norm1 = nn.LayerNorm(dim, eps=1e-6)\n",
    "        self.attn = Attention(\n",
    "                dim,\n",
    "                n_heads=n_heads,\n",
    "                qkv_bias=qkv_bias,\n",
    "                attn_dp=attn_p,\n",
    "                fin_proj_dp=p\n",
    "        )\n",
    "        self.norm2 = nn.LayerNorm(dim, eps=1e-6)\n",
    "        hidden_features = int(dim * mlp_ratio)\n",
    "        self.mlp = MLP(\n",
    "                in_features=dim,\n",
    "                hidden_features=hidden_features,\n",
    "                out_features=dim,\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Run forward pass.\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Shape `(n_samples, n_patches + 1, dim)`.\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Shape `(n_samples, n_patches + 1, dim)`.\n",
    "        \"\"\"\n",
    "        x = x + self.attn(self.norm1(x))\n",
    "        x = x + self.mlp(self.norm2(x))\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65fe4b16",
   "metadata": {},
   "source": [
    "# Vision Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "6b12fd5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VisionTransformer(nn.Module):\n",
    "    \"\"\"\n",
    "    Parameters(img_size,patch_size,in_chans,n_classes,\n",
    "                embed_dims,depth,n_heads,mlp_ratio,qkv_bias,\n",
    "                p,attn_p):\n",
    "    -------\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                img_size=384,patch_size=16,\n",
    "                in_chans=3,n_classes=1000,\n",
    "                embed_dims=768,depth=12,\n",
    "                n_heads=12,mlp_ratio=4,\n",
    "                qkv_bias=True,p=0.0,attn_p=0.0):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.patch_embed=PatchEmbed(img_size=img_size,\n",
    "                                   patch_size=patch_size,\n",
    "                                   in_chans=in_chans,\n",
    "                                   embed_dims=embed_dims)\n",
    "        self.cls_token=nn.Parameter(torch.zeros(1,1,embed_dims))\n",
    "        self.pos_embeds=nn.Parameter(torch.zeros(1,1+self.patch_embed.n_patches,embed_dims))\n",
    "        self.pos_drop=nn.Dropout(p=p)\n",
    "        \n",
    "        self.blocks=nn.ModuleList(\n",
    "                [\n",
    "                    Block(\n",
    "                        dim=embed_dims,\n",
    "                        n_heads=n_heads,\n",
    "                        mlp_ratio=mlp_ratio,\n",
    "                        qkv_bias=qkv_bias,\n",
    "                        p=p,\n",
    "                        attn_p=attn_p\n",
    "                    \n",
    "                    )\n",
    "                    for _ in range(depth)\n",
    "                ]\n",
    "        )\n",
    "        \n",
    "        self.norm=nn.LayerNorm(embed_dims,eps=1e-6)\n",
    "        self.head=nn.Linear(embed_dims,n_classes)\n",
    "        \n",
    "        \n",
    "    def forward(self,x):\n",
    "        \"\"\"\n",
    "        Run the forward pass\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Shape `(n_samples, in_chans, img_size, img_size)`.\n",
    "        Returns\n",
    "        -------\n",
    "        logits : torch.Tensor\n",
    "            Logits over all the classes - `(n_samples, n_classes)`.\n",
    "        \"\"\"\n",
    "        n_samples=x.shape[0]\n",
    "        x=self.patch_embed(x) # will return(n_samples,n_patches,embed_dims)\n",
    "        cls_token=self.cls_token.expand(n_samples,-1,-1) # (n_samples,1,embed_dims)\n",
    "        x=torch.cat((cls_token,x),dim=1)  #(n_samples,1+n_samples,embed_dims)\n",
    "        x=x+ self.pos_embeds\n",
    "        \n",
    "        for block in self.blocks:\n",
    "            x=block(x)\n",
    "            \n",
    "        x=self.norm(x)   # It is always the last dimension which is normalized\n",
    "        \n",
    "        cls_token_final=x[:, 0]  # Just the class token\n",
    "        x=self.head(cls_token_final)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1d362bb",
   "metadata": {},
   "source": [
    "# Testing to check whether our vision transformer is right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "05c9f725",
   "metadata": {},
   "outputs": [],
   "source": [
    "config={\n",
    "        \"img_size\": 384,\n",
    "        \"in_chans\": 3,\n",
    "        \"patch_size\": 16,\n",
    "        \"embed_dims\": 768,\n",
    "        \"depth\": 12,\n",
    "        \"n_heads\": 12,\n",
    "        \"qkv_bias\": True,\n",
    "        \"mlp_ratio\": 4,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "25d7d5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "our_model=VisionTransformer(**config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "097581f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "our_model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "0f3f4c2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install timm -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "e6fbcd21",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "fa91184d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'timm.models.vision_transformer.VisionTransformer'>\n"
     ]
    }
   ],
   "source": [
    "model_name = \"vit_base_patch16_384\"\n",
    "model_official = timm.create_model(model_name, pretrained=True)\n",
    "model_official.eval()\n",
    "print(type(model_official))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "789fcd61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class '__main__.VisionTransformer'>\n"
     ]
    }
   ],
   "source": [
    "print(type(our_model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "0e2a7386",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helpers\n",
    "def get_n_params(module):\n",
    "    return sum(p.numel() for p in module.parameters() if p.requires_grad)\n",
    "\n",
    "def assert_tensors_equal(t1, t2):\n",
    "    a1, a2 = t1.detach().numpy(), t2.detach().numpy()\n",
    "\n",
    "    np.testing.assert_allclose(a1, a2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "87dbfae4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cls_token | cls_token\n",
      "pos_embed | pos_embeds\n",
      "patch_embed.proj.weight | patch_embed.proj.weight\n",
      "patch_embed.proj.bias | patch_embed.proj.bias\n",
      "blocks.0.norm1.weight | blocks.0.norm1.weight\n",
      "blocks.0.norm1.bias | blocks.0.norm1.bias\n",
      "blocks.0.attn.qkv.weight | blocks.0.attn.qkv.weight\n",
      "blocks.0.attn.qkv.bias | blocks.0.attn.qkv.bias\n",
      "blocks.0.attn.proj.weight | blocks.0.attn.fin_proj.weight\n",
      "blocks.0.attn.proj.bias | blocks.0.attn.fin_proj.bias\n",
      "blocks.0.norm2.weight | blocks.0.norm2.weight\n",
      "blocks.0.norm2.bias | blocks.0.norm2.bias\n",
      "blocks.0.mlp.fc1.weight | blocks.0.mlp.fc1.weight\n",
      "blocks.0.mlp.fc1.bias | blocks.0.mlp.fc1.bias\n",
      "blocks.0.mlp.fc2.weight | blocks.0.mlp.fc2.weight\n",
      "blocks.0.mlp.fc2.bias | blocks.0.mlp.fc2.bias\n",
      "blocks.1.norm1.weight | blocks.1.norm1.weight\n",
      "blocks.1.norm1.bias | blocks.1.norm1.bias\n",
      "blocks.1.attn.qkv.weight | blocks.1.attn.qkv.weight\n",
      "blocks.1.attn.qkv.bias | blocks.1.attn.qkv.bias\n",
      "blocks.1.attn.proj.weight | blocks.1.attn.fin_proj.weight\n",
      "blocks.1.attn.proj.bias | blocks.1.attn.fin_proj.bias\n",
      "blocks.1.norm2.weight | blocks.1.norm2.weight\n",
      "blocks.1.norm2.bias | blocks.1.norm2.bias\n",
      "blocks.1.mlp.fc1.weight | blocks.1.mlp.fc1.weight\n",
      "blocks.1.mlp.fc1.bias | blocks.1.mlp.fc1.bias\n",
      "blocks.1.mlp.fc2.weight | blocks.1.mlp.fc2.weight\n",
      "blocks.1.mlp.fc2.bias | blocks.1.mlp.fc2.bias\n",
      "blocks.2.norm1.weight | blocks.2.norm1.weight\n",
      "blocks.2.norm1.bias | blocks.2.norm1.bias\n",
      "blocks.2.attn.qkv.weight | blocks.2.attn.qkv.weight\n",
      "blocks.2.attn.qkv.bias | blocks.2.attn.qkv.bias\n",
      "blocks.2.attn.proj.weight | blocks.2.attn.fin_proj.weight\n",
      "blocks.2.attn.proj.bias | blocks.2.attn.fin_proj.bias\n",
      "blocks.2.norm2.weight | blocks.2.norm2.weight\n",
      "blocks.2.norm2.bias | blocks.2.norm2.bias\n",
      "blocks.2.mlp.fc1.weight | blocks.2.mlp.fc1.weight\n",
      "blocks.2.mlp.fc1.bias | blocks.2.mlp.fc1.bias\n",
      "blocks.2.mlp.fc2.weight | blocks.2.mlp.fc2.weight\n",
      "blocks.2.mlp.fc2.bias | blocks.2.mlp.fc2.bias\n",
      "blocks.3.norm1.weight | blocks.3.norm1.weight\n",
      "blocks.3.norm1.bias | blocks.3.norm1.bias\n",
      "blocks.3.attn.qkv.weight | blocks.3.attn.qkv.weight\n",
      "blocks.3.attn.qkv.bias | blocks.3.attn.qkv.bias\n",
      "blocks.3.attn.proj.weight | blocks.3.attn.fin_proj.weight\n",
      "blocks.3.attn.proj.bias | blocks.3.attn.fin_proj.bias\n",
      "blocks.3.norm2.weight | blocks.3.norm2.weight\n",
      "blocks.3.norm2.bias | blocks.3.norm2.bias\n",
      "blocks.3.mlp.fc1.weight | blocks.3.mlp.fc1.weight\n",
      "blocks.3.mlp.fc1.bias | blocks.3.mlp.fc1.bias\n",
      "blocks.3.mlp.fc2.weight | blocks.3.mlp.fc2.weight\n",
      "blocks.3.mlp.fc2.bias | blocks.3.mlp.fc2.bias\n",
      "blocks.4.norm1.weight | blocks.4.norm1.weight\n",
      "blocks.4.norm1.bias | blocks.4.norm1.bias\n",
      "blocks.4.attn.qkv.weight | blocks.4.attn.qkv.weight\n",
      "blocks.4.attn.qkv.bias | blocks.4.attn.qkv.bias\n",
      "blocks.4.attn.proj.weight | blocks.4.attn.fin_proj.weight\n",
      "blocks.4.attn.proj.bias | blocks.4.attn.fin_proj.bias\n",
      "blocks.4.norm2.weight | blocks.4.norm2.weight\n",
      "blocks.4.norm2.bias | blocks.4.norm2.bias\n",
      "blocks.4.mlp.fc1.weight | blocks.4.mlp.fc1.weight\n",
      "blocks.4.mlp.fc1.bias | blocks.4.mlp.fc1.bias\n",
      "blocks.4.mlp.fc2.weight | blocks.4.mlp.fc2.weight\n",
      "blocks.4.mlp.fc2.bias | blocks.4.mlp.fc2.bias\n",
      "blocks.5.norm1.weight | blocks.5.norm1.weight\n",
      "blocks.5.norm1.bias | blocks.5.norm1.bias\n",
      "blocks.5.attn.qkv.weight | blocks.5.attn.qkv.weight\n",
      "blocks.5.attn.qkv.bias | blocks.5.attn.qkv.bias\n",
      "blocks.5.attn.proj.weight | blocks.5.attn.fin_proj.weight\n",
      "blocks.5.attn.proj.bias | blocks.5.attn.fin_proj.bias\n",
      "blocks.5.norm2.weight | blocks.5.norm2.weight\n",
      "blocks.5.norm2.bias | blocks.5.norm2.bias\n",
      "blocks.5.mlp.fc1.weight | blocks.5.mlp.fc1.weight\n",
      "blocks.5.mlp.fc1.bias | blocks.5.mlp.fc1.bias\n",
      "blocks.5.mlp.fc2.weight | blocks.5.mlp.fc2.weight\n",
      "blocks.5.mlp.fc2.bias | blocks.5.mlp.fc2.bias\n",
      "blocks.6.norm1.weight | blocks.6.norm1.weight\n",
      "blocks.6.norm1.bias | blocks.6.norm1.bias\n",
      "blocks.6.attn.qkv.weight | blocks.6.attn.qkv.weight\n",
      "blocks.6.attn.qkv.bias | blocks.6.attn.qkv.bias\n",
      "blocks.6.attn.proj.weight | blocks.6.attn.fin_proj.weight\n",
      "blocks.6.attn.proj.bias | blocks.6.attn.fin_proj.bias\n",
      "blocks.6.norm2.weight | blocks.6.norm2.weight\n",
      "blocks.6.norm2.bias | blocks.6.norm2.bias\n",
      "blocks.6.mlp.fc1.weight | blocks.6.mlp.fc1.weight\n",
      "blocks.6.mlp.fc1.bias | blocks.6.mlp.fc1.bias\n",
      "blocks.6.mlp.fc2.weight | blocks.6.mlp.fc2.weight\n",
      "blocks.6.mlp.fc2.bias | blocks.6.mlp.fc2.bias\n",
      "blocks.7.norm1.weight | blocks.7.norm1.weight\n",
      "blocks.7.norm1.bias | blocks.7.norm1.bias\n",
      "blocks.7.attn.qkv.weight | blocks.7.attn.qkv.weight\n",
      "blocks.7.attn.qkv.bias | blocks.7.attn.qkv.bias\n",
      "blocks.7.attn.proj.weight | blocks.7.attn.fin_proj.weight\n",
      "blocks.7.attn.proj.bias | blocks.7.attn.fin_proj.bias\n",
      "blocks.7.norm2.weight | blocks.7.norm2.weight\n",
      "blocks.7.norm2.bias | blocks.7.norm2.bias\n",
      "blocks.7.mlp.fc1.weight | blocks.7.mlp.fc1.weight\n",
      "blocks.7.mlp.fc1.bias | blocks.7.mlp.fc1.bias\n",
      "blocks.7.mlp.fc2.weight | blocks.7.mlp.fc2.weight\n",
      "blocks.7.mlp.fc2.bias | blocks.7.mlp.fc2.bias\n",
      "blocks.8.norm1.weight | blocks.8.norm1.weight\n",
      "blocks.8.norm1.bias | blocks.8.norm1.bias\n",
      "blocks.8.attn.qkv.weight | blocks.8.attn.qkv.weight\n",
      "blocks.8.attn.qkv.bias | blocks.8.attn.qkv.bias\n",
      "blocks.8.attn.proj.weight | blocks.8.attn.fin_proj.weight\n",
      "blocks.8.attn.proj.bias | blocks.8.attn.fin_proj.bias\n",
      "blocks.8.norm2.weight | blocks.8.norm2.weight\n",
      "blocks.8.norm2.bias | blocks.8.norm2.bias\n",
      "blocks.8.mlp.fc1.weight | blocks.8.mlp.fc1.weight\n",
      "blocks.8.mlp.fc1.bias | blocks.8.mlp.fc1.bias\n",
      "blocks.8.mlp.fc2.weight | blocks.8.mlp.fc2.weight\n",
      "blocks.8.mlp.fc2.bias | blocks.8.mlp.fc2.bias\n",
      "blocks.9.norm1.weight | blocks.9.norm1.weight\n",
      "blocks.9.norm1.bias | blocks.9.norm1.bias\n",
      "blocks.9.attn.qkv.weight | blocks.9.attn.qkv.weight\n",
      "blocks.9.attn.qkv.bias | blocks.9.attn.qkv.bias\n",
      "blocks.9.attn.proj.weight | blocks.9.attn.fin_proj.weight\n",
      "blocks.9.attn.proj.bias | blocks.9.attn.fin_proj.bias\n",
      "blocks.9.norm2.weight | blocks.9.norm2.weight\n",
      "blocks.9.norm2.bias | blocks.9.norm2.bias\n",
      "blocks.9.mlp.fc1.weight | blocks.9.mlp.fc1.weight\n",
      "blocks.9.mlp.fc1.bias | blocks.9.mlp.fc1.bias\n",
      "blocks.9.mlp.fc2.weight | blocks.9.mlp.fc2.weight\n",
      "blocks.9.mlp.fc2.bias | blocks.9.mlp.fc2.bias\n",
      "blocks.10.norm1.weight | blocks.10.norm1.weight\n",
      "blocks.10.norm1.bias | blocks.10.norm1.bias\n",
      "blocks.10.attn.qkv.weight | blocks.10.attn.qkv.weight\n",
      "blocks.10.attn.qkv.bias | blocks.10.attn.qkv.bias\n",
      "blocks.10.attn.proj.weight | blocks.10.attn.fin_proj.weight\n",
      "blocks.10.attn.proj.bias | blocks.10.attn.fin_proj.bias\n",
      "blocks.10.norm2.weight | blocks.10.norm2.weight\n",
      "blocks.10.norm2.bias | blocks.10.norm2.bias\n",
      "blocks.10.mlp.fc1.weight | blocks.10.mlp.fc1.weight\n",
      "blocks.10.mlp.fc1.bias | blocks.10.mlp.fc1.bias\n",
      "blocks.10.mlp.fc2.weight | blocks.10.mlp.fc2.weight\n",
      "blocks.10.mlp.fc2.bias | blocks.10.mlp.fc2.bias\n",
      "blocks.11.norm1.weight | blocks.11.norm1.weight\n",
      "blocks.11.norm1.bias | blocks.11.norm1.bias\n",
      "blocks.11.attn.qkv.weight | blocks.11.attn.qkv.weight\n",
      "blocks.11.attn.qkv.bias | blocks.11.attn.qkv.bias\n",
      "blocks.11.attn.proj.weight | blocks.11.attn.fin_proj.weight\n",
      "blocks.11.attn.proj.bias | blocks.11.attn.fin_proj.bias\n",
      "blocks.11.norm2.weight | blocks.11.norm2.weight\n",
      "blocks.11.norm2.bias | blocks.11.norm2.bias\n",
      "blocks.11.mlp.fc1.weight | blocks.11.mlp.fc1.weight\n",
      "blocks.11.mlp.fc1.bias | blocks.11.mlp.fc1.bias\n",
      "blocks.11.mlp.fc2.weight | blocks.11.mlp.fc2.weight\n",
      "blocks.11.mlp.fc2.bias | blocks.11.mlp.fc2.bias\n",
      "norm.weight | norm.weight\n",
      "norm.bias | norm.bias\n",
      "head.weight | head.weight\n",
      "head.bias | head.bias\n"
     ]
    }
   ],
   "source": [
    "\n",
    "for (n_o, p_o), (n_c, p_c) in zip(\n",
    "        model_official.named_parameters(), our_model.named_parameters()\n",
    "):\n",
    "    assert p_o.numel() == p_c.numel()\n",
    "    print(f\"{n_o} | {n_c}\")\n",
    "\n",
    "    p_c.data[:] = p_o.data\n",
    "\n",
    "    assert_tensors_equal(p_c.data, p_o.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692bcfbe",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}