transformer_from_scratch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "transformer_from_scratch.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyP3/BVydviVQwbWdoh4gTYp",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/HosseinZaredar/Transformer-from-Scratch/blob/main/transformer_from_scratch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Imports"
      ],
      "metadata": {
        "id": "0ebG4BaQprx3"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T7pZX2CSmi14"
      },
      "outputs": [],
      "source": [
        "import torch as T\n",
        "import torch.nn as nn\n",
        "from torchtext.legacy import data, datasets\n",
        "import torch.optim as optim\n",
        "import time"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Embedding Layer (Word + Positional)\n",
        "\n",
        "Instead of using a sinusoidal positional encoding (as in the original paper), a learnable positional embedding is used."
      ],
      "metadata": {
        "id": "3tdGCSC4LAFP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Embedding(nn.Module):\n",
        "  def __init__(self, vocab_size, max_length, embed_dim, dropout=0.1):\n",
        "    super(Embedding, self).__init__()\n",
        "    self.word_embed = nn.Embedding(vocab_size, embed_dim)\n",
        "    self.pos_embed = nn.Embedding(max_length, embed_dim)\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self, x):\n",
        "    batch_size, seq_length = x.shape\n",
        "    device = T.device('cuda' if T.cuda.is_available() else 'cpu')\n",
        "    positions = T.arange(0, seq_length).expand(\n",
        "        batch_size, seq_length).to(device)\n",
        "    embedding = self.word_embed(x) + self.pos_embed(positions)\n",
        "    return self.dropout(embedding)"
      ],
      "metadata": {
        "id": "L7cA8ytPLFDa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Multi-Head Self-Attention\n",
        "\n",
        "<center>\n",
        "<img src=\"https://raw.githubusercontent.com/HosseinZaredar/Transformer-from-Scratch/main/SelfAttention.png\" width=\"600\" align=\"center\"/>\n",
        "</center>\n"
      ],
      "metadata": {
        "id": "J2BdzpwzqRUG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MHSelfAttention(nn.Module):\n",
        "  def __init__(self, embed_dim, num_heads):\n",
        "    super(MHSelfAttention, self).__init__()\n",
        "    self.embed_dim = embed_dim\n",
        "    self.num_heads = num_heads\n",
        "    self.head_dim = embed_dim // num_heads\n",
        "\n",
        "    assert (self.num_heads*self.head_dim == self.embed_dim), \\\n",
        "    'embed size must be divisible by number of heads'\n",
        "\n",
        "    self.w_queries = nn.Linear(self.embed_dim, self.embed_dim, bias=False)\n",
        "    self.w_keys = nn.Linear(self.embed_dim, self.embed_dim, bias=False)\n",
        "    self.w_values = nn.Linear(self.embed_dim, self.embed_dim, bias=False)\n",
        "    \n",
        "    self.fc_out = nn.Linear(self.head_dim*self.num_heads , self.embed_dim)\n",
        "\n",
        "  def forward(self, x):\n",
        "\n",
        "    # shape of x = [batch_size, sentence_length, embedding_dim]\n",
        "    batch_size = x.shape[0]\n",
        "    sentence_len = x.shape[1]\n",
        "\n",
        "    queries = self.w_queries(x).reshape(\n",
        "        batch_size, sentence_len, self.num_heads, self.head_dim).permute(\n",
        "            0, 2, 1, 3)\n",
        "    \n",
        "    keys = self.w_keys(x).reshape(\n",
        "        batch_size, sentence_len, self.num_heads, self.head_dim).permute(\n",
        "            0, 2, 3, 1)\n",
        "    \n",
        "    \n",
        "    values = self.w_values(x).reshape(\n",
        "        batch_size, sentence_len, self.num_heads, self.head_dim).permute(\n",
        "            0, 2, 1, 3)\n",
        "    \n",
        "    attention_scores = T.einsum('bijk,bikl->bijl', queries, keys)\n",
        "    attention_dist = T.softmax(attention_scores /\n",
        "                               (self.embed_dim ** (1/2)), dim=-1)\n",
        "    attention_out = T.einsum('bijk,bikl->bijl', attention_dist, values)\n",
        "    concatenated_out = attention_out.permute(0, 2, 1, 3).reshape(\n",
        "        batch_size, sentence_len, self.embed_dim)\n",
        "    \n",
        "    return concatenated_out"
      ],
      "metadata": {
        "id": "9dkLjvIeqdQo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Transformer Encoder\n",
        "\n",
        "<center>\n",
        "<img src=\"https://raw.githubusercontent.com/HosseinZaredar/Transformer-from-Scratch/main/Encoder.png\" width=\"200\" align=\"center\"/>\n",
        "</center>"
      ],
      "metadata": {
        "id": "tjtKtvrZ4v22"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerEncoder(nn.Module):\n",
        "  def __init__(self, embed_dim, num_heads, forward_expansion, dropout=0.1):\n",
        "    super(TransformerEncoder, self).__init__()\n",
        "\n",
        "    self.attention = MHSelfAttention(embed_dim, num_heads)\n",
        "    self.norm1 = nn.LayerNorm(embed_dim)\n",
        "    self.norm2 = nn.LayerNorm(embed_dim)\n",
        "\n",
        "    self.feed_forward = nn.Sequential(\n",
        "        nn.Linear(embed_dim, forward_expansion*embed_dim),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(forward_expansion*embed_dim, embed_dim)\n",
        "    )\n",
        "    \n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self, x):\n",
        "    attention_out = self.dropout(self.attention(x))\n",
        "    x = self.norm1(x + attention_out)  \n",
        "    forward_out = self.dropout(self.feed_forward(x))\n",
        "    out = self.norm2(x + forward_out)\n",
        "\n",
        "    return out"
      ],
      "metadata": {
        "id": "NEIb3HI64vFq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### End-to-End Classifier\n",
        "\n",
        "1. An embedding layer\n",
        "2. A single transformer encoder layer\n",
        "3. A fully-connected network as a linear classifier"
      ],
      "metadata": {
        "id": "TJnpDXVjSlhR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Classifier(nn.Module):\n",
        "  def __init__(self, vocab_size, max_length, embed_dim,\n",
        "               num_heads, forward_expansion):\n",
        "      super(Classifier, self).__init__()\n",
        "\n",
        "      self.embedder = Embedding(vocab_size, max_length, embed_dim)\n",
        "      self.encoder = TransformerEncoder(embed_dim, num_heads, forward_expansion)\n",
        "      self.fc = nn.Linear(embed_dim, 1)\n",
        "\n",
        "  def forward(self, x):\n",
        "    embedding = self.embedder(x)\n",
        "    encoding = self.encoder(embedding)\n",
        "    compact_encoding = encoding.max(dim=1)[0]\n",
        "    out = self.fc(compact_encoding)\n",
        "    return out"
      ],
      "metadata": {
        "id": "x9qLIrNhSsbk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load and Preprocess IMDb Dataset"
      ],
      "metadata": {
        "id": "kN2DGXTsILSb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "TEXT = data.Field(\n",
        "    tokenize = 'spacy', tokenizer_language = 'en_core_web_sm',batch_first=True)\n",
        "LABEL = data.LabelField(dtype = T.float)"
      ],
      "metadata": {
        "id": "uRdBH6BmIKrd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data, test_data = datasets.IMDB.splits(TEXT, LABEL)"
      ],
      "metadata": {
        "id": "sxBDZRlZMTLT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data, valid_data = train_data.split()"
      ],
      "metadata": {
        "id": "2WAxuSGRMhPh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "MAX_VOCAB_SIZE = 25000\n",
        "\n",
        "TEXT.build_vocab(train_data, max_size = MAX_VOCAB_SIZE)\n",
        "LABEL.build_vocab(train_data)"
      ],
      "metadata": {
        "id": "rPKDLy_NNaCm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "BATCH_SIZE = 4\n",
        "device = T.device('cuda' if T.cuda.is_available() else 'cpu')\n",
        "\n",
        "train_iterator, valid_iterator, test_iterator = data.BucketIterator.splits(\n",
        "    (train_data, valid_data, test_data), \n",
        "    batch_size = BATCH_SIZE,\n",
        "    device = device)"
      ],
      "metadata": {
        "id": "iQTSI_VANdHq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Initializing The Model"
      ],
      "metadata": {
        "id": "z_K9btG7SB_Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "EMBED_DIM = 200\n",
        "NUM_HEADS = 8\n",
        "FORWARD_EXPANSION = 3\n",
        "MAX_LENGTH = 512\n",
        "VOCAB_SIZE = len(TEXT.vocab)\n",
        "\n",
        "classifier = Classifier(\n",
        "    VOCAB_SIZE, MAX_LENGTH, EMBED_DIM, NUM_HEADS, FORWARD_EXPANSION)\n",
        "device = T.device('cuda' if T.cuda.is_available() else 'cpu')\n",
        "classifier.to(device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WWSyihBxN4Sj",
        "outputId": "c68c0d2f-2dce-47e7-f9cc-352db11cebff"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Classifier(\n",
              "  (embedder): Embedding(\n",
              "    (word_embed): Embedding(25002, 200)\n",
              "    (pos_embed): Embedding(512, 200)\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              "  (encoder): TransformerEncoder(\n",
              "    (attention): MHSelfAttention(\n",
              "      (w_queries): Linear(in_features=200, out_features=200, bias=False)\n",
              "      (w_keys): Linear(in_features=200, out_features=200, bias=False)\n",
              "      (w_values): Linear(in_features=200, out_features=200, bias=False)\n",
              "      (fc_out): Linear(in_features=200, out_features=200, bias=True)\n",
              "    )\n",
              "    (norm1): LayerNorm((200,), eps=1e-05, elementwise_affine=True)\n",
              "    (norm2): LayerNorm((200,), eps=1e-05, elementwise_affine=True)\n",
              "    (feed_forward): Sequential(\n",
              "      (0): Linear(in_features=200, out_features=600, bias=True)\n",
              "      (1): ReLU()\n",
              "      (2): Linear(in_features=600, out_features=200, bias=True)\n",
              "    )\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              "  (fc): Linear(in_features=200, out_features=1, bias=True)\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training"
      ],
      "metadata": {
        "id": "w08nIXt-_OpH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "optimizer = optim.SGD(classifier.parameters(), lr=1e-3)"
      ],
      "metadata": {
        "id": "nmML4fPjSIZk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "criterion = nn.BCEWithLogitsLoss()\n",
        "device = T.device('cuda' if T.cuda.is_available() else 'cpu')\n",
        "criterion.to(device);"
      ],
      "metadata": {
        "id": "M8KtZobUYJHA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def binary_accuracy(preds, y):\n",
        "    rounded_preds = T.round(T.sigmoid(preds))\n",
        "    correct = (rounded_preds == y).float()\n",
        "    acc = correct.sum() / len(correct)\n",
        "    return acc"
      ],
      "metadata": {
        "id": "Jjs2XuaDbs8k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def epoch_time(start_time, end_time):\n",
        "    elapsed_time = end_time - start_time\n",
        "    elapsed_mins = int(elapsed_time / 60)\n",
        "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
        "    return elapsed_mins, elapsed_secs"
      ],
      "metadata": {
        "id": "uGvO_3PbAq7q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(model, iterator, optimizer, criterion):\n",
        "    \n",
        "    epoch_loss = 0\n",
        "    epoch_acc = 0\n",
        "    \n",
        "    model.train()\n",
        "    \n",
        "    for batch in iterator:\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        input = batch.text\n",
        "        if input.shape[1] > MAX_LENGTH:\n",
        "          input = input[:, :MAX_LENGTH]\n",
        "                \n",
        "        predictions = model(input).squeeze(1)\n",
        "        \n",
        "        loss = criterion(predictions, batch.label)\n",
        "        \n",
        "        acc = binary_accuracy(predictions, batch.label)\n",
        "        \n",
        "        loss.backward()\n",
        "        \n",
        "        optimizer.step()\n",
        "        \n",
        "        epoch_loss += loss.item()\n",
        "        epoch_acc += acc.item()\n",
        "        \n",
        "    return epoch_loss / len(iterator), epoch_acc / len(iterator)"
      ],
      "metadata": {
        "id": "T_0AQJgEYMjL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(model, iterator, criterion):\n",
        "    epoch_loss = 0\n",
        "    epoch_acc = 0\n",
        "    model.eval()\n",
        "\n",
        "    with T.no_grad():\n",
        "        for batch in iterator:\n",
        "\n",
        "            input = batch.text\n",
        "            if input.shape[1] > MAX_LENGTH:\n",
        "              input = input[:, :MAX_LENGTH]\n",
        "                    \n",
        "            predictions = model(input).squeeze(1)\n",
        "            loss = criterion(predictions, batch.label)\n",
        "            acc = binary_accuracy(predictions, batch.label)\n",
        "            epoch_loss += loss.item()\n",
        "            epoch_acc += acc.item()\n",
        "        \n",
        "    return epoch_loss / len(iterator), epoch_acc / len(iterator)"
      ],
      "metadata": {
        "id": "oRV4SusuuGek"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "N_EPOCHS = 10\n",
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(N_EPOCHS):\n",
        "\n",
        "    start_time = time.time()\n",
        "\n",
        "    train_loss, train_acc = train(\n",
        "        classifier, train_iterator,optimizer, criterion)\n",
        "    valid_loss, valid_acc = evaluate(classifier, valid_iterator, criterion)\n",
        "\n",
        "    end_time = time.time()\n",
        "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
        "\n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        T.save(classifier.state_dict(), 'sent-classifier.pt')\n",
        "\n",
        "    print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
        "    print(f'\\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')\n",
        "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eo4Kw1rpb7un",
        "outputId": "98416562-2a45-41ea-e4a9-9444942131f3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 01 | Epoch Time: 1m 2s\n",
            "\tTrain Loss: 0.730 | Train Acc: 49.83%\n",
            "\t Val. Loss: 0.702 |  Val. Acc: 50.51%\n",
            "Epoch: 02 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.719 | Train Acc: 51.80%\n",
            "\t Val. Loss: 0.688 |  Val. Acc: 53.43%\n",
            "Epoch: 03 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.704 | Train Acc: 53.22%\n",
            "\t Val. Loss: 0.689 |  Val. Acc: 53.61%\n",
            "Epoch: 04 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.695 | Train Acc: 55.47%\n",
            "\t Val. Loss: 0.683 |  Val. Acc: 55.73%\n",
            "Epoch: 05 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.684 | Train Acc: 57.75%\n",
            "\t Val. Loss: 0.684 |  Val. Acc: 56.09%\n",
            "Epoch: 06 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.670 | Train Acc: 59.65%\n",
            "\t Val. Loss: 0.641 |  Val. Acc: 62.57%\n",
            "Epoch: 07 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.648 | Train Acc: 62.01%\n",
            "\t Val. Loss: 0.618 |  Val. Acc: 65.43%\n",
            "Epoch: 08 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.627 | Train Acc: 64.58%\n",
            "\t Val. Loss: 0.595 |  Val. Acc: 67.84%\n",
            "Epoch: 09 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.602 | Train Acc: 67.27%\n",
            "\t Val. Loss: 0.586 |  Val. Acc: 68.76%\n",
            "Epoch: 10 | Epoch Time: 1m 1s\n",
            "\tTrain Loss: 0.579 | Train Acc: 69.30%\n",
            "\t Val. Loss: 0.575 |  Val. Acc: 70.03%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluation"
      ],
      "metadata": {
        "id": "7_4MS275uFBO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "classifier.load_state_dict(T.load('sent-classifier.pt'))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QYM_vp1dA490",
        "outputId": "f54bc069-faa8-4a14-c74b-5b9385dece69"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "metadata": {},
          "execution_count": 45
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_loss, test_acc = evaluate(classifier, test_iterator, criterion)\n",
        "print(f'Test Loss: {test_loss:.3f} |  Test Acc: {test_acc*100:.2f}%')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "N1a03U6Cuc3e",
        "outputId": "e920559a-6652-467a-b9b9-4fd39d40e098"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Test Loss: 0.578 |  Test Acc: 69.58%\n"
          ]
        }
      ]
    }
  ]
}