14_Transformer.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "08_pytorch_paper_replicating_exercises.ipynb",
   "provenance": [],
   "collapsed_sections": [],
   "toc_visible": true,
   "authorship_tag": "ABX9TyOhoCjGZZxrecbm76R8UJZn",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3 (ipykernel)",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard"
 },
 "cells": [
  {
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import random_split, DataLoader\n",
    "from torchinfo import summary\n",
    "from torchvision import datasets, transforms, models\n",
    "\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(DEVICE)"
   ],
   "metadata": {
    "id": "Y5H5P8EjCNGK",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "4214da9e-f3a8-43e6-a48c-1a33f44a9be9",
    "ExecuteTime": {
     "end_time": "2024-10-14T15:06:46.034872Z",
     "start_time": "2024-10-14T15:06:46.007834Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "execution_count": 27
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Implementatation"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T17:13:24.173148Z",
     "start_time": "2024-10-14T17:13:24.094357Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, input_dim, embedding_dim):\n",
    "        super().__init__()\n",
    "        self.input_dim = input_dim\n",
    "        \n",
    "        self.Q = nn.Linear(input_dim, embedding_dim)\n",
    "        self.K = nn.Linear(input_dim, embedding_dim)\n",
    "        self.V = nn.Linear(input_dim, embedding_dim)\n",
    "    \n",
    "    def forward(self, x_q, x_k, x_v, mask=None):\n",
    "        q = self.Q(x_q)\n",
    "        k = self.K(x_k)\n",
    "        v = self.V(x_v)\n",
    "        \n",
    "        a = (q @ k.permute(0, 2, 1)) / (self.input_dim ** 0.5)\n",
    "        if mask is not None:\n",
    "            a[:, mask] = -1e10\n",
    "        x = torch.softmax(a, dim=2) @ v\n",
    "        return x\n",
    "\n",
    "\n",
    "class MultiheadSelfAttention(nn.Module):\n",
    "    def __init__(self, input_dim, n_heads):\n",
    "        super().__init__()\n",
    "        self.attention_heads = nn.ModuleList([SelfAttention(input_dim, input_dim // n_heads) for _ in range(n_heads)])\n",
    "        self.linear = nn.Linear(input_dim, input_dim)\n",
    "    \n",
    "    def forward(self, x_q, x_k, x_v, mask=None):\n",
    "        x = torch.concat([attention_head(x_q, x_k, x_v, mask) for attention_head in self.attention_heads], dim=-1)\n",
    "        x = self.linear(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "x_batch = torch.randn(4, 100, 512)\n",
    "mask = torch.ones(100, 100).tril() == 0\n",
    "\n",
    "print(MultiheadSelfAttention(512, 16)(x_batch, x_batch, x_batch, mask=mask).shape)\n",
    "print(nn.MultiheadAttention(512, 16, batch_first=True)(x_batch, x_batch, x_batch, need_weights=False, attn_mask=mask)[0].shape)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 100, 512])\n",
      "torch.Size([4, 100, 512])\n"
     ]
    }
   ],
   "execution_count": 147
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T17:13:24.319372Z",
     "start_time": "2024-10-14T17:13:24.314088Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class MSABlock(nn.Module):\n",
    "    def __init__(self, input_dim, n_heads, dropout, torch_msa=True):\n",
    "        super().__init__()\n",
    "        self.torch_msa = torch_msa\n",
    "        \n",
    "        self.msa = nn.MultiheadAttention(input_dim, n_heads, batch_first=True) if torch_msa \\\n",
    "            else MultiheadSelfAttention(input_dim, n_heads)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x_q, x_k, x_v, mask=None):\n",
    "        x = self.msa(x_q, x_k, x_v, need_weights=False, attn_mask=mask)[0] if self.torch_msa \\\n",
    "            else self.msa(x_q, x_k, x_v, mask=mask)\n",
    "        x = self.dropout(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class MLPBlock(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, dropout):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.linear1 = nn.Linear(input_dim, mlp_dim)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.linear2 = nn.Linear(mlp_dim, input_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.linear1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.linear2(x)\n",
    "        x = self.dropout(x)\n",
    "        return x"
   ],
   "outputs": [],
   "execution_count": 148
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T17:13:24.519591Z",
     "start_time": "2024-10-14T17:13:24.514416Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, n_heads, dropout):\n",
    "        super().__init__()\n",
    "        self.msa_block = MSABlock(input_dim, n_heads, dropout)\n",
    "        self.layer_norm1 = nn.LayerNorm(input_dim)\n",
    "        self.mlp_block = MLPBlock(input_dim, mlp_dim, dropout)\n",
    "        self.layer_norm2 = nn.LayerNorm(input_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.msa_block(x, x, x) + x\n",
    "        x = self.layer_norm1(x)\n",
    "        x = self.mlp_block(x) + x\n",
    "        x = self.layer_norm2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, n_heads, dropout, n_layers):\n",
    "        super().__init__()\n",
    "        self.encoder_layers = nn.Sequential(\n",
    "            *[EncoderBlock(input_dim, mlp_dim, n_heads, dropout) for _ in range(n_layers)]\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.encoder_layers(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# summary(Encoder(512, 2048, 8, 0.1, 6), input_size=(1, 500, 512), device='cpu', col_names=['output_size', 'num_params', 'mult_adds'], col_width=15)"
   ],
   "outputs": [],
   "execution_count": 149
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T17:13:24.968481Z",
     "start_time": "2024-10-14T17:13:24.915110Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, n_heads, dropout):\n",
    "        super().__init__()\n",
    "        self.masked_msa_block = MSABlock(input_dim, n_heads, dropout)\n",
    "        self.layer_norm1 = nn.LayerNorm(input_dim)\n",
    "        self.enc_msa_block = MSABlock(input_dim, n_heads, dropout)\n",
    "        self.layer_norm2 = nn.LayerNorm(input_dim)\n",
    "        self.mlp_block = MLPBlock(input_dim, mlp_dim, dropout)\n",
    "        self.layer_norm3 = nn.LayerNorm(input_dim)\n",
    "    \n",
    "    def forward(self, x, enc_out, mask=None):\n",
    "        x = self.masked_msa_block(x, x, x) + x\n",
    "        x = self.layer_norm1(x)\n",
    "        x = self.enc_msa_block(x, enc_out, enc_out, mask=mask) + x # mask!\n",
    "        x = self.layer_norm2(x)\n",
    "        x = self.mlp_block(x) + x\n",
    "        x = self.layer_norm3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "x_batch = torch.randn(4, 100, 512)\n",
    "print(DecoderBlock(512, 2048, 8, 0.1)(x_batch, x_batch).shape)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 100, 512])\n"
     ]
    }
   ],
   "execution_count": 150
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T17:13:25.590413Z",
     "start_time": "2024-10-14T17:13:25.581061Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, seq_len, word_dim, embedding_dim, mlp_dim, n_heads, n_layers, n_classes, dropout):\n",
    "        super().__init__()\n",
    "        self.seq_len = seq_len\n",
    "        self.embedding_dim = embedding_dim\n",
    "        \n",
    "        self.input_embedding = nn.Linear(word_dim, embedding_dim)\n",
    "        self.output_embedding = nn.Linear(word_dim, embedding_dim)\n",
    "        self.positional_encoding = self.create_positional_encoding()\n",
    "        \n",
    "        self.encoder_layers = Encoder(embedding_dim, mlp_dim, n_heads, dropout, n_layers)\n",
    "        self.decoder_layers = nn.ModuleList([\n",
    "            DecoderBlock(embedding_dim, mlp_dim, n_heads, dropout) for _ in range(n_layers)\n",
    "        ])\n",
    "        \n",
    "        self.classification_head = nn.Sequential(\n",
    "            nn.LayerNorm(embedding_dim),\n",
    "            nn.Linear(embedding_dim, n_classes)\n",
    "        )\n",
    "    \n",
    "    def create_positional_encoding(self):\n",
    "        pos = torch.arange(self.seq_len).unsqueeze(1)\n",
    "        denominator = 10000 ** (torch.arange(0, self.embedding_dim, 2) / self.embedding_dim)\n",
    "        \n",
    "        pos_encoding = torch.zeros(seq_len, self.embedding_dim)\n",
    "        pos_encoding[:, 0::2] = torch.sin(pos / denominator)\n",
    "        pos_encoding[:, 1::2] = torch.cos(pos / denominator)\n",
    "        return pos_encoding\n",
    "    \n",
    "    def create_attention_mask(self):\n",
    "        mask = torch.ones(self.seq_len, self.seq_len).tril() == 0\n",
    "        return mask\n",
    "    \n",
    "    def forward(self, input_seq, output_seq):\n",
    "        input_embedding = self.input_embedding(input_seq) + self.positional_encoding\n",
    "        output_embedding = self.output_embedding(output_seq) + self.positional_encoding\n",
    "        \n",
    "        encoder_out = self.encoder_layers(input_embedding)\n",
    "        \n",
    "        mask = self.create_attention_mask()\n",
    "        decoder_out = output_embedding\n",
    "        for decoder_layer in self.decoder_layers:\n",
    "            decoder_layer(decoder_out, encoder_out, mask)\n",
    "        logits = self.classification_head(decoder_out)\n",
    "        \n",
    "        return logits\n",
    "\n",
    "\n",
    "# summary(Transformer(100, 1000, 512, 2048, 8, 6, 1000, 0.1), input_size=[(1, 100, 1000), (1, 100, 1000)], device='cpu', col_names=['output_size', 'num_params', 'mult_adds'], col_width=13)"
   ],
   "outputs": [],
   "execution_count": 151
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T17:13:28.733208Z",
     "start_time": "2024-10-14T17:13:27.924441Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# seq_len, word_dim, embedding_dim, mlp_dim, n_heads, n_layers, n_classes, dropout\n",
    "\n",
    "seq_len = 100\n",
    "word_dim = 1000\n",
    "embedding_dim = 512\n",
    "mlp_dim = embedding_dim*4\n",
    "n_heads = 8\n",
    "n_layers = 6\n",
    "dropout = 0.1\n",
    "\n",
    "transformer_model = Transformer(seq_len, word_dim, embedding_dim, mlp_dim, n_heads, n_layers, 1000, dropout)\n",
    "summary(transformer_model, input_size=[(1, 100, 1000), (1, 100, 1000)], device='cpu', col_names=['output_size', 'num_params', 'mult_adds'], depth=3)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==================================================================================================================================\n",
       "Layer (type:depth-idx)                                  Output Shape              Param #                   Mult-Adds\n",
       "==================================================================================================================================\n",
       "Transformer                                             [1, 100, 1000]            --                        --\n",
       "├─Linear: 1-1                                           [1, 100, 512]             512,512                   512,512\n",
       "├─Linear: 1-2                                           [1, 100, 512]             512,512                   512,512\n",
       "├─Encoder: 1-3                                          [1, 100, 512]             --                        --\n",
       "│    └─Sequential: 2-1                                  [1, 100, 512]             --                        --\n",
       "│    │    └─EncoderBlock: 3-1                           [1, 100, 512]             3,152,384                 3,152,384\n",
       "│    │    └─EncoderBlock: 3-2                           [1, 100, 512]             3,152,384                 3,152,384\n",
       "│    │    └─EncoderBlock: 3-3                           [1, 100, 512]             3,152,384                 3,152,384\n",
       "│    │    └─EncoderBlock: 3-4                           [1, 100, 512]             3,152,384                 3,152,384\n",
       "│    │    └─EncoderBlock: 3-5                           [1, 100, 512]             3,152,384                 3,152,384\n",
       "│    │    └─EncoderBlock: 3-6                           [1, 100, 512]             3,152,384                 3,152,384\n",
       "├─ModuleList: 1-4                                       --                        --                        --\n",
       "│    └─DecoderBlock: 2-2                                [1, 100, 512]             --                        --\n",
       "│    │    └─MSABlock: 3-7                               [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-8                              [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MSABlock: 3-9                               [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-10                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MLPBlock: 3-11                              [1, 100, 512]             2,099,712                 2,099,712\n",
       "│    │    └─LayerNorm: 3-12                             [1, 100, 512]             1,024                     1,024\n",
       "│    └─DecoderBlock: 2-3                                [1, 100, 512]             --                        --\n",
       "│    │    └─MSABlock: 3-13                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-14                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MSABlock: 3-15                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-16                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MLPBlock: 3-17                              [1, 100, 512]             2,099,712                 2,099,712\n",
       "│    │    └─LayerNorm: 3-18                             [1, 100, 512]             1,024                     1,024\n",
       "│    └─DecoderBlock: 2-4                                [1, 100, 512]             --                        --\n",
       "│    │    └─MSABlock: 3-19                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-20                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MSABlock: 3-21                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-22                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MLPBlock: 3-23                              [1, 100, 512]             2,099,712                 2,099,712\n",
       "│    │    └─LayerNorm: 3-24                             [1, 100, 512]             1,024                     1,024\n",
       "│    └─DecoderBlock: 2-5                                [1, 100, 512]             --                        --\n",
       "│    │    └─MSABlock: 3-25                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-26                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MSABlock: 3-27                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-28                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MLPBlock: 3-29                              [1, 100, 512]             2,099,712                 2,099,712\n",
       "│    │    └─LayerNorm: 3-30                             [1, 100, 512]             1,024                     1,024\n",
       "│    └─DecoderBlock: 2-6                                [1, 100, 512]             --                        --\n",
       "│    │    └─MSABlock: 3-31                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-32                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MSABlock: 3-33                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-34                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MLPBlock: 3-35                              [1, 100, 512]             2,099,712                 2,099,712\n",
       "│    │    └─LayerNorm: 3-36                             [1, 100, 512]             1,024                     1,024\n",
       "│    └─DecoderBlock: 2-7                                [1, 100, 512]             --                        --\n",
       "│    │    └─MSABlock: 3-37                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-38                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MSABlock: 3-39                              [1, 100, 512]             1,050,624                 1,050,624\n",
       "│    │    └─LayerNorm: 3-40                             [1, 100, 512]             1,024                     1,024\n",
       "│    │    └─MLPBlock: 3-41                              [1, 100, 512]             2,099,712                 2,099,712\n",
       "│    │    └─LayerNorm: 3-42                             [1, 100, 512]             1,024                     1,024\n",
       "├─Sequential: 1-5                                       [1, 100, 1000]            --                        --\n",
       "│    └─LayerNorm: 2-8                                   [1, 100, 512]             1,024                     1,024\n",
       "│    └─Linear: 2-9                                      [1, 100, 1000]            513,000                   513,000\n",
       "==================================================================================================================================\n",
       "Total params: 45,677,544\n",
       "Trainable params: 45,677,544\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 45.68\n",
       "==================================================================================================================================\n",
       "Input size (MB): 0.80\n",
       "Forward/backward pass size (MB): 68.38\n",
       "Params size (MB): 182.71\n",
       "Estimated Total Size (MB): 251.89\n",
       "=================================================================================================================================="
      ]
     },
     "execution_count": 152,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 152
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T13:20:16.241309Z",
     "start_time": "2024-10-14T13:20:15.503311Z"
    }
   },
   "cell_type": "code",
   "source": [
    "transformer_torch_model = nn.Transformer(embedding_dim, n_heads, n_layers, n_layers, mlp_dim, dropout)\n",
    "summary(transformer_torch_model, input_size=[(1, 100, 512), (1, 100, 512)], device='cpu', col_names=['output_size', 'num_params', 'mult_adds'], depth=3)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=============================================================================================================================\n",
       "Layer (type:depth-idx)                             Output Shape              Param #                   Mult-Adds\n",
       "=============================================================================================================================\n",
       "Transformer                                        [1, 100, 512]             --                        --\n",
       "├─TransformerEncoder: 1-1                          [1, 100, 512]             --                        --\n",
       "│    └─ModuleList: 2-1                             --                        --                        --\n",
       "│    │    └─TransformerEncoderLayer: 3-1           [1, 100, 512]             3,152,384                 2,101,760\n",
       "│    │    └─TransformerEncoderLayer: 3-2           [1, 100, 512]             3,152,384                 2,101,760\n",
       "│    │    └─TransformerEncoderLayer: 3-3           [1, 100, 512]             3,152,384                 2,101,760\n",
       "│    │    └─TransformerEncoderLayer: 3-4           [1, 100, 512]             3,152,384                 2,101,760\n",
       "│    │    └─TransformerEncoderLayer: 3-5           [1, 100, 512]             3,152,384                 2,101,760\n",
       "│    │    └─TransformerEncoderLayer: 3-6           [1, 100, 512]             3,152,384                 2,101,760\n",
       "│    └─LayerNorm: 2-2                              [1, 100, 512]             1,024                     1,024\n",
       "├─TransformerDecoder: 1-2                          [1, 100, 512]             --                        --\n",
       "│    └─ModuleList: 2-3                             --                        --                        --\n",
       "│    │    └─TransformerDecoderLayer: 3-7           [1, 100, 512]             4,204,032                 2,102,784\n",
       "│    │    └─TransformerDecoderLayer: 3-8           [1, 100, 512]             4,204,032                 2,102,784\n",
       "│    │    └─TransformerDecoderLayer: 3-9           [1, 100, 512]             4,204,032                 2,102,784\n",
       "│    │    └─TransformerDecoderLayer: 3-10          [1, 100, 512]             4,204,032                 2,102,784\n",
       "│    │    └─TransformerDecoderLayer: 3-11          [1, 100, 512]             4,204,032                 2,102,784\n",
       "│    │    └─TransformerDecoderLayer: 3-12          [1, 100, 512]             4,204,032                 2,102,784\n",
       "│    └─LayerNorm: 2-4                              [1, 100, 512]             1,024                     1,024\n",
       "=============================================================================================================================\n",
       "Total params: 44,140,544\n",
       "Trainable params: 44,140,544\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 25.23\n",
       "=============================================================================================================================\n",
       "Input size (MB): 0.41\n",
       "Forward/backward pass size (MB): 37.68\n",
       "Params size (MB): 100.92\n",
       "Estimated Total Size (MB): 139.01\n",
       "============================================================================================================================="
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  }
 ]
}