3.2-Transformer-intro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: implement Multi Head Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Multi Head Attention](../labs/images/multi-head-attention.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Multi-Head Attention module.\n",
    "\n",
    "    This module performs multi-head attention on the input tensor `x`. It applies linear transformations\n",
    "    to the input tensor to obtain query, key, and value tensors. Then, it calculates attention weights\n",
    "    and applies attention to the value tensor. Finally, it applies another linear transformation to\n",
    "    obtain the output tensor.\n",
    "\n",
    "    Args:\n",
    "        in_channels (int): Number of input channels.\n",
    "        hidden_channels (int): Number of hidden channels.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        out_channels (int, optional): Number of output channels. If not provided, it is set to `in_channels`.\n",
    "        bias (bool, optional): Whether to include bias in linear transformations. Default is False.\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Output tensor of shape (batch_size, seq_len, in_channels).\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, in_channels, hidden_channels, num_heads, out_channels=None, bias=False\n",
    "    ):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.hidden_channels = hidden_channels\n",
    "        self.in_channels = in_channels\n",
    "\n",
    "        self.W_q = nn.Linear(in_channels, num_heads * hidden_channels, bias=bias)\n",
    "        self.W_k = nn.Linear(in_channels, num_heads * hidden_channels, bias=bias)\n",
    "        self.W_v = nn.Linear(in_channels, num_heads * hidden_channels, bias=bias)\n",
    "        self.W_o = nn.Linear(\n",
    "            num_heads * hidden_channels,\n",
    "            in_channels if out_channels is None else out_channels,\n",
    "        )\n",
    "\n",
    "        self.scale = 1.0 / math.sqrt(hidden_channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x.shape = (batch_size, seq_len, in_channels)\n",
    "\n",
    "        B = x.shape[0]\n",
    "        N = x.shape[1]\n",
    "\n",
    "        \"\"\"\n",
    "        insert the missing codes here\n",
    "        \"\"\"\n",
    "        Q = ...\n",
    "        K = ...\n",
    "        V = ...\n",
    "\n",
    "        # reshape\n",
    "        Q = Q.view(B, self.num_heads, N, self.hidden_channels)\n",
    "        K = K.view(B, self.num_heads, N, self.hidden_channels)\n",
    "        V = V.view(B, self.num_heads, N, self.hidden_channels)\n",
    "\n",
    "        # calculate attention\n",
    "        \"\"\"\n",
    "        Hint: you need to transpose the last two dimensions of K before you can multiply it with Q to get the attention logits.\n",
    "        \"\"\"\n",
    "        attn_logits = ...\n",
    "        attn_logits *= self.scale\n",
    "        attn_weights = F.softmax(attn_logits, dim=-1)\n",
    "        # attn_weights.shape: (B, num_heads, N, N)\n",
    "\n",
    "        # apply attention\n",
    "        mha_output = ...\n",
    "        # mha_output.shape: (B, num_heads, N, hidden_channels)\n",
    "\n",
    "        # reshape\n",
    "        mha_output = mha_output.view(B, N, self.hidden_channels * self.num_heads)\n",
    "        # mha_output.shape: (B, N, hidden_channels * num_heads)\n",
    "\n",
    "        mha_output = self.W_o(mha_output)\n",
    "        # mha_output.shape: (B, N, in_channels)\n",
    "\n",
    "        return mha_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mha = MultiHeadAttention(32, 32, 4)\n",
    "x = torch.rand(2, 5, 32)\n",
    "output = mha(x)\n",
    "print(f\"Input shape: {x.shape}\")\n",
    "print(f\"Output shape: {output.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composing the Transformer Block"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Transformer encoder](../labs/images/transformer-encoder.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardNetwork(nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels):\n",
    "        super(FeedForwardNetwork, self).__init__()\n",
    "        self.ff = nn.Sequential(\n",
    "            nn.Linear(in_channels, hidden_channels),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_channels, out_channels),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.ff(x)\n",
    "\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        hidden_channels = d_model // num_heads\n",
    "        self.mha = MultiHeadAttention(d_model, hidden_channels, num_heads)\n",
    "        self.ln1 = nn.LayerNorm(d_model)\n",
    "        self.ln2 = nn.LayerNorm(d_model)\n",
    "        self.ffn = FeedForwardNetwork(d_model, d_model * 4, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        mha_output = self.mha(x)\n",
    "        x = self.ln1(x + mha_output)\n",
    "        ffn_output = self.ffn(x)\n",
    "        x = self.ln2(x + ffn_output)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_block = TransformerBlock(32, 4)\n",
    "x = torch.rand(2, 5, 32)\n",
    "output = transformer_block(x)\n",
    "\n",
    "print(f\"Input shape: {x.shape}\")\n",
    "print(f\"Output shape: {output.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: define the Transformer Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, num_layers=1):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        self.layers = nn.ModuleList(\n",
    "            [TransformerBlock(d_model, num_heads) for _ in range(num_layers)]\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x.shape = (batch_size, seq_len, d_model)\n",
    "        \"\"\"\n",
    "        insert your code here\n",
    "        \"\"\"\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: build a Vision transformer\n",
    "We build a slight variation from the model below, namely we stick to the original Transformer Block architecture and just a single linear layer to project the output of the transformer to the number of classes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![ViT](../labs/images/vit-overview.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VisionTransformer(nn.Module):\n",
    "    def __init__(\n",
    "        self, in_channels, d_model, num_heads, num_classes, n_patches, num_layers=1\n",
    "    ):\n",
    "        super(VisionTransformer, self).__init__()\n",
    "        # positional embeddings\n",
    "        self.pos_embedding = nn.Parameter(torch.randn(1, n_patches + 1, d_model) * 0.02)\n",
    "        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))\n",
    "        self.encoder = TransformerEncoder(d_model, num_heads, num_layers)\n",
    "        self.fc_in = nn.Linear(in_channels, d_model)\n",
    "        self.fc_out = nn.Linear(d_model, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        fill in the blanks here\n",
    "        \"\"\"\n",
    "        # x.shape = (batch_size, n_patches, in_channels)\n",
    "        B, N, C = x.shape\n",
    "        x = self.fc_in(x)  # (B, N, C) -> (B, N, d_model)\n",
    "        cls_tokens = self.cls_token.expand(\n",
    "            B, -1, -1\n",
    "        )  # (1, 1, d_model) -> (B, 1, d_model)\n",
    "\n",
    "        # concatenate the cls token with the input\n",
    "        x = ...\n",
    "\n",
    "        # add the positonal embeddings\n",
    "        x += self.pos_embedding  # (B, N+1, d_model)\n",
    "        \n",
    "        # apply the encoder\n",
    "        x = ...  # (B, N+1, d_model)\n",
    "        # extract the cls token\n",
    "        x = ...  # (B, d_model)\n",
    "        \n",
    "        x = self.fc_out(x)  # (B, num_classes)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Vision Transformer on MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Load MNIST dataset\n",
    "train_dataset = datasets.MNIST(\n",
    "    root=\"./data\", train=True, download=True, transform=transform\n",
    ")\n",
    "test_dataset = datasets.MNIST(\n",
    "    root=\"./data\", train=False, download=True, transform=transform\n",
    ")\n",
    "\n",
    "# reduce the size of the dataset for faster training\n",
    "train_dataset.data = train_dataset.data[:1000]\n",
    "train_dataset.targets = train_dataset.targets[:1000]\n",
    "\n",
    "test_dataset.data = test_dataset.data[:100]\n",
    "test_dataset.targets = test_dataset.targets[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def patchify_img(img, patch_size=7):\n",
    "    c, h, w = img.shape\n",
    "    assert h / patch_size == h // patch_size, \"height must be divisible by patch size\"\n",
    "    assert w / patch_size == w // patch_size, \"width must be divisible by patch size\"\n",
    "    patches = img.reshape(c, h // patch_size, patch_size, -1, patch_size)\n",
    "    patches = patches.permute(0, 1, 3, 2, 4).reshape(-1, patch_size, patch_size)\n",
    "    return patches\n",
    "\n",
    "\n",
    "class PatchifyDataset(Dataset):\n",
    "    def __init__(self, dataset, patch_size):\n",
    "        self.dataset = dataset\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img, label = self.dataset[idx]\n",
    "        patches = patchify_img(img, self.patch_size)\n",
    "        # flatten the last dim\n",
    "        patches = patches.view(patches.size(0), -1)\n",
    "        return patches, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "n_epochs = 50\n",
    "patch_size = 7\n",
    "num_layers = 1\n",
    "num_heads = 4\n",
    "seed = 42\n",
    "\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "train_dataset_patched = PatchifyDataset(train_dataset, patch_size=patch_size)\n",
    "test_dataset_patched = PatchifyDataset(test_dataset, patch_size=patch_size)\n",
    "\n",
    "train_loader = DataLoader(train_dataset_patched, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset_patched, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using {device}\")\n",
    "\n",
    "model = VisionTransformer(\n",
    "    in_channels=patch_size**2,\n",
    "    d_model=32,\n",
    "    num_heads=num_heads,\n",
    "    num_classes=10,\n",
    "    n_patches=16,\n",
    "    num_layers=num_layers,\n",
    ").to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    model.train()\n",
    "    running_train_loss = 0.0\n",
    "    for x, y in train_loader:\n",
    "        x, y = x.to(device), y.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = model(x)\n",
    "        loss = criterion(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_train_loss += loss.item()\n",
    "    running_train_loss /= len(train_loader)\n",
    "\n",
    "    model.eval()\n",
    "    running_val_loss = 0.0\n",
    "    with torch.no_grad():\n",
    "        for x, y in test_loader:\n",
    "            x, y = x.to(device), y.to(device)\n",
    "            y_pred = model(x)\n",
    "            val_loss = criterion(y_pred, y)\n",
    "            running_val_loss += val_loss.item()\n",
    "        running_val_loss /= len(test_loader)\n",
    "\n",
    "    print(\n",
    "        f\"Epoch {epoch}, Train loss: {running_train_loss}, Val. loss: {running_val_loss}\"\n",
    "    )\n",
    "    train_losses.append(running_train_loss)\n",
    "    val_losses.append(running_val_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(train_losses, label=\"train loss\")\n",
    "plt.plot(val_losses, label=\"val loss\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus: apply the model on the full test set, visualize some results and compute the accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further reading/watching"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. [GPT from scratch by Andrey Karpathy](https://www.youtube.com/watch?v=kCc8FmEb1nY), his other videos are also extremely informative"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spacet",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}