assignments/assignment3/Transformer_Captioning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This mounts your Google Drive to the Colab VM.\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "# TODO: Enter the foldername in your Drive where you have saved the unzipped\n",
    "# assignment folder, e.g. 'cs231n/assignments/assignment3/'\n",
    "FOLDERNAME = None\n",
    "assert FOLDERNAME is not None, \"[!] Enter the foldername.\"\n",
    "\n",
    "# Now that we've mounted your Drive, this ensures that\n",
    "# the Python interpreter of the Colab VM can load\n",
    "# python files from within it.\n",
    "import sys\n",
    "sys.path.append('/content/drive/My Drive/{}'.format(FOLDERNAME))\n",
    "\n",
    "# This downloads the COCO dataset to your Drive\n",
    "# if it doesn't already exist.\n",
    "%cd /content/drive/My\\ Drive/$FOLDERNAME/cs231n/datasets/\n",
    "!bash get_datasets.sh\n",
    "%cd /content/drive/My\\ Drive/$FOLDERNAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-title"
    ]
   },
   "source": [
    "# Image Captioning with Transformers\n",
    "You have now implemented a vanilla RNN and for the task of image captioning. In this notebook you will implement key pieces of a transformer decoder to accomplish the same task.\n",
    "\n",
    "**NOTE:** This notebook will be primarily written in PyTorch rather than NumPy, unlike the RNN notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "# Setup cell.\n",
    "import time, os, json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from cs231n.gradient_check import eval_numerical_gradient, eval_numerical_gradient_array\n",
    "from cs231n.transformer_layers import *\n",
    "from cs231n.captioning_solver_transformer import CaptioningSolverTransformer\n",
    "from cs231n.classifiers.transformer import CaptioningTransformer\n",
    "from cs231n.coco_utils import load_coco_data, sample_coco_minibatch, decode_captions\n",
    "from cs231n.image_utils import image_from_url\n",
    "\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (10.0, 8.0) # Set default size of plots.\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "plt.rcParams['image.cmap'] = 'gray'\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "def rel_error(x, y):\n",
    "    \"\"\" returns relative error \"\"\"\n",
    "    return np.max(np.abs(x - y) / (np.maximum(1e-8, np.abs(x) + np.abs(y))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# COCO Dataset\n",
    "As in the previous notebooks, we will use the COCO dataset for captioning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "base dir  E:\\Mo\\Univercity\\internship\\cs231n\\assignments\\assignment3\\cs231n\\datasets/coco_captioning\n",
      "train_captions <class 'numpy.ndarray'> (400135, 17) int32\n",
      "train_image_idxs <class 'numpy.ndarray'> (400135,) int32\n",
      "val_captions <class 'numpy.ndarray'> (195954, 17) int32\n",
      "val_image_idxs <class 'numpy.ndarray'> (195954,) int32\n",
      "train_features <class 'numpy.ndarray'> (82783, 512) float32\n",
      "val_features <class 'numpy.ndarray'> (40504, 512) float32\n",
      "idx_to_word <class 'list'> 1004\n",
      "word_to_idx <class 'dict'> 1004\n",
      "train_urls <class 'numpy.ndarray'> (82783,) <U63\n",
      "val_urls <class 'numpy.ndarray'> (40504,) <U63\n"
     ]
    }
   ],
   "source": [
    "# Load COCO data from disk into a dictionary.\n",
    "data = load_coco_data(pca_features=True)\n",
    "\n",
    "# Print out all the keys and values from the data dictionary.\n",
    "for k, v in data.items():\n",
    "    if type(v) == np.ndarray:\n",
    "        print(k, type(v), v.shape, v.dtype)\n",
    "    else:\n",
    "        print(k, type(v), len(v))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer\n",
    "As you have seen, RNNs are incredibly powerful but often slow to train. Further, RNNs struggle to encode long-range dependencies (though LSTMs are one way of mitigating the issue). In 2017, Vaswani et al introduced the Transformer in their paper [\"Attention Is All You Need\"](https://arxiv.org/abs/1706.03762) to a) introduce parallelism and b) allow models to learn long-range dependencies. The paper not only led to famous models like BERT and GPT in the natural language processing community, but also an explosion of interest across fields, including vision. While here we introduce the model in the context of image captioning, the idea of attention itself is much more general.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KqPMDm4F9m0v"
   },
   "source": [
    "# Transformer: Multi-Headed Attention\n",
    "\n",
    "### Dot-Product Attention\n",
    "\n",
    "Recall that attention can be viewed as an operation on a query $q\\in\\mathbb{R}^d$, a set of value vectors $\\{v_1,\\dots,v_n\\}, v_i\\in\\mathbb{R}^d$, and a set of key vectors $\\{k_1,\\dots,k_n\\}, k_i \\in \\mathbb{R}^d$, specified as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align}\n",
    "c = \\sum_{i=1}^{n} v_i \\alpha_i \\\\ \\alpha_i = \\frac{\\exp(k_i^\\top q)}{\\sum_{j=1}^{n} \\exp(k_j^\\top q)} \\\\\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $\\alpha_i$ are frequently called the \"attention weights\", and the output $c\\in\\mathbb{R}^d$ is a correspondingly weighted average over the value vectors.\n",
    "\n",
    "### Self-Attention\n",
    "In Transformers, we perform self-attention, which means that the values, keys and query are derived from the input $X \\in \\mathbb{R}^{\\ell \\times d}$, where $\\ell$ is our sequence length. Specifically, we learn parameter matrices $V,K,Q \\in \\mathbb{R}^{d\\times d}$ to map our input $X$ as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align}\n",
    "v_i = Vx_i\\ \\ i \\in \\{1,\\dots,\\ell\\}\\\\\n",
    "k_i = Kx_i\\ \\ i \\in \\{1,\\dots,\\ell\\}\\\\\n",
    "q_i = Qx_i\\ \\ i \\in \\{1,\\dots,\\ell\\}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Headed Scaled Dot-Product Attention\n",
    "In the case of multi-headed attention, we learn a parameter matrix for each head, which gives the model more expressivity to attend to different parts of the input. Let $h$ be number of heads, and $Y_i$ be the attention output of head $i$. Thus we learn individual matrices $Q_i$, $K_i$ and $V_i$. To keep our overall computation the same as the single-headed case, we choose $Q_i \\in \\mathbb{R}^{d\\times d/h}$, $K_i \\in \\mathbb{R}^{d\\times d/h}$ and $V_i \\in \\mathbb{R}^{d\\times d/h}$. Adding in a scaling term $\\frac{1}{\\sqrt{d/h}}$ to our simple dot-product attention above, we have"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation} \\label{qkv_eqn}\n",
    "Y_i = \\text{softmax}\\bigg(\\frac{(XQ_i)(XK_i)^\\top}{\\sqrt{d/h}}\\bigg)(XV_i)\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $Y_i\\in\\mathbb{R}^{\\ell \\times d/h}$, where $\\ell$ is our sequence length.\n",
    "\n",
    "In our implementation, we then apply dropout here (though in practice it could be used at any step):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$Y_i = \\text{dropout}(Y_i)$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, then the output of the self-attention is a linear transformation of the concatenation of the heads:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "Y = [Y_1;\\dots;Y_h]A\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "were $A \\in\\mathbb{R}^{d\\times d}$ and $[Y_1;\\dots;Y_h]\\in\\mathbb{R}^{\\ell \\times d}$.\n",
    "\n",
    "Implement multi-headed scaled dot-product attention in the `MultiHeadAttention` class in the file `cs231n/transformer_layers.py`. The code below will check your implementation. The relative error should be less than 1e-3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(231)\n",
    "\n",
    "# Choose dimensions such that they are all unique for easier debugging:\n",
    "# Specifically, the following values correspond to N=1, H=2, T=3, E//H=4, and E=8.\n",
    "batch_size = 1\n",
    "sequence_length = 3\n",
    "embed_dim = 8\n",
    "attn = MultiHeadAttention(embed_dim, num_heads=2)\n",
    "\n",
    "# Self-attention.\n",
    "data = torch.randn(batch_size, sequence_length, embed_dim)\n",
    "self_attn_output = attn(query=data, key=data, value=data)\n",
    "\n",
    "# Masked self-attention.\n",
    "mask = torch.randn(sequence_length, sequence_length) < 0.5\n",
    "masked_self_attn_output = attn(query=data, key=data, value=data, attn_mask=mask)\n",
    "\n",
    "# Attention using two inputs.\n",
    "other_data = torch.randn(batch_size, sequence_length, embed_dim)\n",
    "attn_output = attn(query=data, key=other_data, value=other_data)\n",
    "\n",
    "expected_self_attn_output = np.asarray([[\n",
    "[-0.2494,  0.1396,  0.4323, -0.2411, -0.1547,  0.2329, -0.1936,\n",
    "          -0.1444],\n",
    "         [-0.1997,  0.1746,  0.7377, -0.3549, -0.2657,  0.2693, -0.2541,\n",
    "          -0.2476],\n",
    "         [-0.0625,  0.1503,  0.7572, -0.3974, -0.1681,  0.2168, -0.2478,\n",
    "          -0.3038]]])\n",
    "\n",
    "expected_masked_self_attn_output = np.asarray([[\n",
    "[-0.1347,  0.1934,  0.8628, -0.4903, -0.2614,  0.2798, -0.2586,\n",
    "          -0.3019],\n",
    "         [-0.1013,  0.3111,  0.5783, -0.3248, -0.3842,  0.1482, -0.3628,\n",
    "          -0.1496],\n",
    "         [-0.2071,  0.1669,  0.7097, -0.3152, -0.3136,  0.2520, -0.2774,\n",
    "          -0.2208]]])\n",
    "\n",
    "expected_attn_output = np.asarray([[\n",
    "[-0.1980,  0.4083,  0.1968, -0.3477,  0.0321,  0.4258, -0.8972,\n",
    "          -0.2744],\n",
    "         [-0.1603,  0.4155,  0.2295, -0.3485, -0.0341,  0.3929, -0.8248,\n",
    "          -0.2767],\n",
    "         [-0.0908,  0.4113,  0.3017, -0.3539, -0.1020,  0.3784, -0.7189,\n",
    "          -0.2912]]])\n",
    "\n",
    "print('self_attn_output error: ', rel_error(expected_self_attn_output, self_attn_output.detach().numpy()))\n",
    "print('masked_self_attn_output error: ', rel_error(expected_masked_self_attn_output, masked_self_attn_output.detach().numpy()))\n",
    "print('attn_output error: ', rel_error(expected_attn_output, attn_output.detach().numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tcDBRnqL9m0w"
   },
   "source": [
    "# Positional Encoding\n",
    "\n",
    "While transformers are able to easily attend to any part of their input, the attention mechanism has no concept of token order. However, for many tasks (especially natural language processing), relative token order is very important. To recover this, the authors add a positional encoding to the embeddings of individual word tokens.\n",
    "\n",
    "Let us define a matrix $P \\in \\mathbb{R}^{l\\times d}$, where $P_{ij} = $"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{cases}\n",
    "\\text{sin}\\left(i \\cdot 10000^{-\\frac{j}{d}}\\right) & \\text{if j is even} \\\\\n",
    "\\text{cos}\\left(i \\cdot 10000^{-\\frac{(j-1)}{d}}\\right) & \\text{otherwise} \\\\\n",
    "\\end{cases}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rather than directly passing an input $X \\in \\mathbb{R}^{l\\times d}$ to our network, we instead pass $X + P$.\n",
    "\n",
    "Implement this layer in `PositionalEncoding` in `cs231n/transformer_layers.py`. Once you are done, run the following to perform a simple test of your implementation. You should see errors on the order of `e-3` or less."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(231)\n",
    "\n",
    "batch_size = 1\n",
    "sequence_length = 2\n",
    "embed_dim = 6\n",
    "data = torch.randn(batch_size, sequence_length, embed_dim)\n",
    "\n",
    "pos_encoder = PositionalEncoding(embed_dim)\n",
    "output = pos_encoder(data)\n",
    "\n",
    "expected_pe_output = np.asarray([[[-1.2340,  1.1127,  1.6978, -0.0865, -0.0000,  1.2728],\n",
    "                                  [ 0.9028, -0.4781,  0.5535,  0.8133,  1.2644,  1.7034]]])\n",
    "\n",
    "print('pe_output error: ', rel_error(expected_pe_output, output.detach().numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-inline"
    ]
   },
   "source": [
    "# Inline Question 1\n",
    "\n",
    "Several key design decisions were made in designing the scaled dot product attention we introduced above. Explain why the following choices were beneficial:\n",
    "1. Using multiple attention heads as opposed to one.\n",
    "2. Dividing by $\\sqrt{d/h}$ before applying the softmax function. Recall that $d$ is the feature dimension and $h$ is the number of heads.\n",
    "3. Adding a linear transformation to the output of the attention operation. What would happen if we were to stack attention operations directly?\n",
    "\n",
    "Only one or two sentences per choice is necessary, but be sure to be specific in addressing what would have happened without each given implementation detail, why such a situation would be suboptimal, and how the proposed implementation improves the situation.\n",
    "\n",
    "**Your Answer:** \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overfit Transformer Captioning Model on Small Data\n",
    "Run the following to overfit the Transformer-based captioning model on the same small dataset as we used for the RNN previously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(231)\n",
    "np.random.seed(231)\n",
    "\n",
    "data = load_coco_data(max_train=50)\n",
    "\n",
    "transformer = CaptioningTransformer(\n",
    "          word_to_idx=data['word_to_idx'],\n",
    "          input_dim=data['train_features'].shape[1],\n",
    "          wordvec_dim=256,\n",
    "          num_heads=2,\n",
    "          num_layers=2,\n",
    "          max_length=30\n",
    "        )\n",
    "\n",
    "\n",
    "transformer_solver = CaptioningSolverTransformer(transformer, data, idx_to_word=data['idx_to_word'],\n",
    "           num_epochs=100,\n",
    "           batch_size=25,\n",
    "           learning_rate=0.001,\n",
    "           verbose=True, print_every=10,\n",
    "         )\n",
    "\n",
    "transformer_solver.train()\n",
    "\n",
    "# Plot the training losses.\n",
    "plt.plot(transformer_solver.loss_history)\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training loss history')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print final training loss. You should see a final loss of less than 0.03."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "transformer_final_training_loss",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print('Final loss: ', transformer_solver.loss_history[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer Sampling at Test Time\n",
    "The sampling code has been written for you. You can simply run the following to compare with the previous results with the RNN. As before the training results should be much better than the validation set results, given how little data we trained on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# If you get an error, the URL just no longer exists, so don't worry!\n",
    "# You can re-sample as many times as you want.\n",
    "for split in ['train', 'val']:\n",
    "    minibatch = sample_coco_minibatch(data, split=split, batch_size=2)\n",
    "    gt_captions, features, urls = minibatch\n",
    "    gt_captions = decode_captions(gt_captions, data['idx_to_word'])\n",
    "\n",
    "    sample_captions = transformer.sample(features, max_length=30)\n",
    "    sample_captions = decode_captions(sample_captions, data['idx_to_word'])\n",
    "\n",
    "    for gt_caption, sample_caption, url in zip(gt_captions, sample_captions, urls):\n",
    "        img = image_from_url(url)\n",
    "        # Skip missing URLs.\n",
    "        if img is None: continue\n",
    "        plt.imshow(img)            \n",
    "        plt.title('%s\\n%s\\nGT:%s' % (split, sample_caption, gt_caption))\n",
    "        plt.axis('off')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}