8.1 - Transformer Tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer activity\n",
    "In this activity, we will build a transformer to do a regression task based on discrete sequential inputs (very similar to the RNA project)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "DEVICE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup the data\n",
    "\n",
    "We will build a custom dataset for this dataset.\n",
    "\n",
    "For the input data:\n",
    "- `0` represents padding\n",
    "- `1` represents \"Up\"\n",
    "- `2` represents \"Down\"\n",
    "- `3` represents \"Left\"\n",
    "- `4` represents \"Right\"\n",
    "\n",
    "Outputs are also set to 0 for padded indices. Both inputs and outputs have a maximum sequence length of 32.\n",
    "\n",
    "We should also make an attention mask that is `True` when the input is 0 and `False` otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For Colab\n",
    "# !git clone https://github.com/trevor-yu-087/syde-599-f23-tutorial.git\n",
    "# DATA_PATH = \"syde-599-f23-tutorial/data/transformer-data.pkl\"\n",
    "# DATA_PATH = \"./data/transformer-data.pkl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(DATA_PATH, \"rb\") as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split into train and test datasets\n",
    "X_train, X_test, y_train, y_test = train_test_split(data[\"inputs\"], data[\"targets\"], test_size=10_000, random_state=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create custom torch map-style dataset by implementing __len__ and __getitem__\n",
    "# Example: https://pytorch.org/tutorials/beginner/data_loading_tutorial.html#dataset-class\n",
    "\n",
    "class AntHillDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, inputs, targets):\n",
    "        \"\"\"\n",
    "        inputs: list of np.array, input move sequences encoded as integers. \n",
    "        targets: list of np.array, output height values.\n",
    "        \"\"\"\n",
    "        self.inputs = inputs\n",
    "        self.targets = targets\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.targets)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"Retreive one item of the dataset at specified index.\n",
    "        Parameters:\n",
    "        -----------\n",
    "        idx: int\n",
    "\n",
    "        Returns: \n",
    "        --------\n",
    "        (input, target, padding_mask): Tuple of torch.LongTensor, torch.FloatTensor, torch.BoolTensor.\n",
    "        All have shape (1, 32). padding_mask is True when that position is a padding element.\n",
    "        \"\"\"\n",
    "        input = self.inputs[idx]\n",
    "        target = self.targets[idx]\n",
    "        padding_mask = input == 0  # True when 0, do not attend to padded items\n",
    "        \n",
    "        input = torch.LongTensor(input)\n",
    "        target = torch.FloatTensor(target)\n",
    "        padding_mask = torch.BoolTensor(padding_mask)\n",
    "        \n",
    "        return input, target, padding_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = AntHillDataset(X_train, y_train)\n",
    "test_ds = AntHillDataset(X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To improve dataloader performance, pinning memory during dataloading can help. You will notice performance improvements after the first epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=True, pin_memory=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_ds, batch_size=128, shuffle=False, pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sinusoidal Position Encodings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Complete the formula for the position encoding (the argument to the sine and cosine). The odd dimension indices get sine features, while the even dimension indices get cosine features. $D$ is the embedding dimension, and $N$ is the sequence length.\n",
    "\n",
    "$\\mathbf{p}[k, 2i] = sin(k / 10000^{2i / D}); \\ i = 1,...,D//2; \\ k = 1,...,N$\n",
    "\n",
    "$\\mathbf{p}[k, 2i+1] = cos(k / 10000^{2i / D}); \\ i = 1,...,D//2; \\ k = 1,...,N$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "class SinusoidalPositionalEmbedding(nn.Embedding):\n",
    "    \"\"\"This module produces sinusoidal positional embeddings of any length.\"\"\"\n",
    "    def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None) -> None:\n",
    "        super().__init__(num_positions, embedding_dim)\n",
    "        self.weight = self._init_weight(self.weight)\n",
    "\n",
    "    @staticmethod\n",
    "    def _init_weight(out: nn.Parameter) -> nn.Parameter:\n",
    "        \"\"\"\n",
    "        Interleaved sine and cosine position embeddings\n",
    "        \"\"\"\n",
    "        out.requires_grad = False\n",
    "        out.detach_()\n",
    "        \n",
    "        N, D = out.shape\n",
    "\n",
    "        ## TODO: Create a N x D//2 array of position encodings (argument to the sine/cosine)\n",
    "        position_enc = None\n",
    "        #####\n",
    "\n",
    "        out[:, 0::2] = torch.FloatTensor(np.sin(position_enc))  # Even indices get sin\n",
    "        out[:, 1::2] = torch.FloatTensor(np.cos(position_enc))  # Odd indices get cos\n",
    "        return out\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def forward(self, input_ids_shape: torch.Size, past_key_values_length: int = 0) -> torch.Tensor:\n",
    "        \"\"\"`input_ids_shape` is expected to be [bsz x seqlen].\"\"\"\n",
    "        bsz, seq_len = input_ids_shape[:2]\n",
    "        positions = torch.arange(\n",
    "            past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device\n",
    "        )\n",
    "        return super().forward(positions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Run this cell to test your implementation\n",
    "input_ids = torch.zeros(64, 256, 30)\n",
    "\n",
    "spe = SinusoidalPositionalEmbedding(num_positions=256, embedding_dim=32)\n",
    "# Note that the forward takes the input shape as the argument\n",
    "out = spe(input_ids.shape)\n",
    "\n",
    "# Plot some dimensions to compare\n",
    "import matplotlib.pyplot as plt\n",
    "inds_to_plot = [5, 10, 15, 16, 20, 25]\n",
    "for i in inds_to_plot:\n",
    "    plt.plot(out[:, i], label=f\"d = {i}\")\n",
    "plt.xlabel(\"Position\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer model\n",
    "Transformer models typically have 3 stages:\n",
    "1. Input stage: Embedding layers to process inputs\n",
    "2. Processing stage: Transformer blocks\n",
    "3. Output stage: Process sequential output for specific task (e.g. classification head, regression head)\n",
    "\n",
    "For this model, we'll use the `nn.Embedding` module to create vector embeddings for our input space, and sinusoidal position encodings, though a learned position embedding could also be used too.\n",
    "\n",
    "For the transformer blocks, we'll use a bi-directional transformer encoder. We'll use the `GELU` activation, norm-first architecture, and batch first tensor shapes.\n",
    "\n",
    "Note for tensor inputs:\n",
    "- N = 32 (sequence length)\n",
    "- D = 128 (hidden dimension)\n",
    "- H = 4 (number of heads)\n",
    "- D_FFN = 2 * D\n",
    "- Input shape: (B, N, D)\n",
    "- Output shape: (B, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AntHillTransformer(nn.Module):\n",
    "    def __init__(self, d_model=128, n_head=4, n_layers=3):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.n_head = n_head\n",
    "        self.n_layers = n_layers\n",
    "\n",
    "        # Input embedding layers\n",
    "        # TODO: Create input embedding and position encoding layers\n",
    "        self.input_embeddings = None\n",
    "        self.position_encoding = None\n",
    "        ###\n",
    "        self.input_layer_norm = nn.LayerNorm(d_model)\n",
    "\n",
    "        # Transformer layers\n",
    "        # TODO: Create encoder layer and transformer encoder stack\n",
    "        encoder_layer = nn.TransformerEncoderLayer(\n",
    "            \n",
    "        )\n",
    "        self.transformer_encoder = nn.TransformerEncoder()\n",
    "        ###\n",
    "\n",
    "        # Output layers\n",
    "        self.output_norm = nn.LayerNorm(d_model)  # With pre-layer norm transformer, we should do one more before output\n",
    "        # TODO: Create output linear layer\n",
    "        self.output_layer = None\n",
    "        ###\n",
    "\n",
    "    def forward(self, input_ids, padding_mask):\n",
    "        \"\"\"\n",
    "        Parameters:\n",
    "        -----------\n",
    "        inputs: torch.LongTensor of shape (B, N, D)\n",
    "        padding_mask: torch.BoolTensor of shape (B, N)\n",
    "        \n",
    "        Returns:\n",
    "        --------\n",
    "        outputs: torch.FlaotTensor of shape (B, N)\n",
    "        \"\"\"\n",
    "        # TODO: Fill in forward layer\n",
    "        # Inputs\n",
    "        inputs = None\n",
    "        inputs = self.input_layer_norm(inputs)\n",
    "\n",
    "        # Transformer\n",
    "        outputs = None\n",
    "\n",
    "        # Output layers\n",
    "        outputs = None\n",
    "        ###\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AntHillTransformer()\n",
    "\n",
    "(inputs, targets, mask) = next(iter(test_loader))\n",
    "outputs = model(inputs, mask)\n",
    "outputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "m = mask[i, :].detach().cpu().numpy()\n",
    "tgt = targets[i, :].detach().cpu().numpy()\n",
    "y = outputs[i, :].detach().cpu().numpy()\n",
    "\n",
    "plt.plot(tgt[~m], label=\"Target\")\n",
    "plt.plot(y[~m], label=\"Prediction\")\n",
    "plt.legend()\n",
    "plt.title(\"Model prediction before training\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss function\n",
    "\n",
    "We need our loss function to only compute the loss over the non-masked examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Outputs should be shape (B, N) = (128, 32)\n",
    "def masked_mse(outputs, targets, mask):\n",
    "    # TODO: Compute loss over non-masked elements\n",
    "    loss = None\n",
    "    ###\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = masked_mse(outputs, targets, mask)\n",
    "loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup training loop\n",
    "Note we modify the training loop from MNIST to include the mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, loss_fn, optimizer, epoch=-1):\n",
    "    \"\"\"\n",
    "    Trains a model for one epoch (one pass through the entire training data).\n",
    "\n",
    "    :param model: PyTorch model\n",
    "    :param train_loader: PyTorch Dataloader for training data\n",
    "    :param loss_fn: PyTorch loss function\n",
    "    :param optimizer: PyTorch optimizer, initialized with model parameters\n",
    "    :kwarg epoch: Integer epoch to use when printing loss and accuracy\n",
    "    :returns: Accuracy score\n",
    "    \"\"\"\n",
    "    total_loss = 0\n",
    "    all_predictions = []\n",
    "    all_targets = []\n",
    "    loss_history = []\n",
    "\n",
    "    model = model.to(DEVICE)\n",
    "    model.train()  # Set model in training mode\n",
    "    for i, (inputs, targets, mask) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        inputs, targets, mask = inputs.to(DEVICE), targets.to(DEVICE), mask.to(DEVICE)\n",
    "        outputs = model(inputs, mask)\n",
    "        loss = loss_fn(outputs, targets, mask)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Track some values to compute statistics\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # Save loss every 100 batches\n",
    "        if (i % 100 == 0) and (i > 0):\n",
    "            running_loss = total_loss / (i + 1)\n",
    "            loss_history.append(running_loss)\n",
    "            # print(f\"Epoch {epoch + 1}, batch {i + 1}: loss = {running_loss:.2f}\")\n",
    "\n",
    "    final_loss = total_loss / len(train_loader)\n",
    "    # Print average loss and accuracy\n",
    "    print(f\"Epoch {epoch + 1} done. Average train loss = {final_loss:.2f}\")\n",
    "    return final_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, test_loader, loss_fn, epoch=-1):\n",
    "    \"\"\"\n",
    "    Tests a model for one epoch of test data.\n",
    "\n",
    "    Note:\n",
    "        In testing and evaluation, we do not perform gradient descent optimization, so steps 2, 5, and 6 are not needed.\n",
    "        For performance, we also tell torch not to track gradients by using the `with torch.no_grad()` context.\n",
    "\n",
    "    :param model: PyTorch model\n",
    "    :param test_loader: PyTorch Dataloader for test data\n",
    "    :param loss_fn: PyTorch loss function\n",
    "    :kwarg epoch: Integer epoch to use when printing loss and accuracy\n",
    "\n",
    "    :returns: Accuracy score\n",
    "    \"\"\"\n",
    "    total_loss = 0\n",
    "    all_predictions = []\n",
    "    all_targets = []\n",
    "    model = model.to(DEVICE)\n",
    "    model.eval()  # Set model in evaluation mode\n",
    "    for i, (inputs, targets, mask) in enumerate(test_loader):\n",
    "        inputs, targets, mask = inputs.to(DEVICE), targets.to(DEVICE), mask.to(DEVICE)\n",
    "        with torch.no_grad():\n",
    "            outputs = model(inputs, mask)\n",
    "            loss = loss_fn(outputs, targets, mask)\n",
    "\n",
    "            # Track some values to compute statistics\n",
    "            total_loss += loss.item()\n",
    "\n",
    "    final_loss = total_loss / len(test_loader)\n",
    "    # Print average loss and accuracy\n",
    "    print(f\"Epoch {epoch + 1} done. Average test loss = {final_loss:.2f}\")\n",
    "    return final_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test the implementation\n",
    "Train a model for 5 epochs with both custom and PyTorch optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 1e-3\n",
    "NUM_EPOCHS = 5\n",
    "torch.manual_seed(0)\n",
    "model = AntHillTransformer()\n",
    "optimizer = torch.optim.Adam(model.parameters(), LEARNING_RATE)\n",
    "loss_fn = masked_mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_losses = []\n",
    "test_losses = []\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    train_loss = train(model, train_loader, loss_fn, optimizer, epoch)\n",
    "    test_loss= test(model, test_loader, loss_fn, epoch)\n",
    "\n",
    "    train_losses.append(train_loss)\n",
    "    test_losses.append(test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View an example from the test dataset\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    (inputs, targets, mask) = next(iter(test_loader))\n",
    "    outputs = model(inputs, mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "m = mask[i, :].cpu().numpy()\n",
    "tgt = targets[i, :].cpu().numpy()\n",
    "y = outputs[i, :].cpu().numpy()\n",
    "\n",
    "plt.plot(tgt[~m], label=\"Target\")\n",
    "plt.plot(y[~m], label=\"Prediction\")\n",
    "plt.legend()\n",
    "plt.title(\"Model prediction after training\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "datasci",
   "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": 2
}