1_Architecture.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preparation &  Architecture Design\n",
    "\n",
    "In this notebook, we will explain how to load and pre-process data from the [COCO dataset](http://cocodataset.org/#home). We will also design a CNN-RNN model for automatically generating image captions.\n",
    "\n",
    "The implementation of CNN encoder and RNN decoder are in the **model.py** file.  \n",
    "\n",
    "Outline of this notebook:\n",
    "- [Step 1](#step1): Explore the Data Loader\n",
    "- [Step 2](#step2): Use the Data Loader to Obtain Batches\n",
    "- [Step 3](#step3): Experiment with the CNN Encoder\n",
    "- [Step 4](#step4): Implement the RNN Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step1'></a>\n",
    "## Step 1: Write the Data Loader\n",
    "\n",
    "We wrote a [data loader](http://pytorch.org/docs/master/data.html#torch.utils.data.DataLoader) to load the COCO dataset in batches. In the code cell below, We will initialize the data loader by using the `get_loader` function in **data_loader.py**.  \n",
    "\n",
    "\n",
    "The `get_loader` function takes as input a number of arguments that can be explored in **data_loader.py**.  Take the time to explore these arguments now by opening **data_loader.py**.  Most of the arguments must be left at their default values Here some important parameters:\n",
    "1. **`transform`** - an [image transform](http://pytorch.org/docs/master/torchvision/transforms.html) specifying how to pre-process the images and convert them to PyTorch tensors before using them as input to the CNN encoder.  We defined transforms in `transform_train` variable.\n",
    "2. **`mode`** - one of `'train'` (loads the training data in batches) or `'test'` (for the test data). We will say that the data loader is in training or test mode, respectively.  We keep the data loader in training mode in this notebook by setting `mode='train'`.\n",
    "3. **`batch_size`** - determines the batch size.  When training the model, this is number of image-caption pairs used to amend the model weights in each training step.\n",
    "4. **`vocab_threshold`** - **the total number of times that a word must appear in the in the training captions before it is used as part of the vocabulary**.  Words that have fewer than `vocab_threshold` occurrences in the training captions are considered unknown words. \n",
    "5. **`vocab_from_file`** - a Boolean that decides whether to load the vocabulary from file.  \n",
    "\n",
    "We will describe the `vocab_threshold` and `vocab_from_file` arguments in more detail soon. For now, run the code cell below. It may take a couple of minutes to run!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": false
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /home/masoud/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n",
      "[nltk_data] Downloading package punkt to /home/masoud/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "from pycocotools.coco import COCO\n",
    "import nltk\n",
    "from data_loader import get_loader\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.utils.data as data\n",
    "from torchvision import transforms\n",
    "\n",
    "nltk.download(\"punkt\")\n",
    "\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.61s)\n",
      "creating index...\n",
      "index created!\n",
      "[0/414113] Tokenizing captions...\n",
      "[100000/414113] Tokenizing captions...\n",
      "[200000/414113] Tokenizing captions...\n",
      "[300000/414113] Tokenizing captions...\n",
      "[400000/414113] Tokenizing captions...\n",
      "loading annotations into memory...\n",
      "Done (t=0.64s)\n",
      "creating index...\n",
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 414113/414113 [00:45<00:00, 9162.46it/s]\n"
     ]
    }
   ],
   "source": [
    "# Define a transform to pre-process the training images.\n",
    "transform_train = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(256),  # smaller edge of image resized to 256\n",
    "        transforms.RandomCrop(224),  # get 224x224 crop from random location\n",
    "        transforms.RandomHorizontalFlip(),  # horizontally flip image with probability=0.5\n",
    "        transforms.ToTensor(),  # convert the PIL Image to a tensor\n",
    "        transforms.Normalize(\n",
    "            (0.485, 0.456, 0.406),  # normalize image for pre-trained model\n",
    "            (0.229, 0.224, 0.225),\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Set the minimum word count threshold.\n",
    "vocab_threshold = 5\n",
    "\n",
    "# Specify the batch size.\n",
    "batch_size = 10\n",
    "\n",
    "# Path to cocoapi dir\n",
    "cocoapi_dir = r\"path/to/cocoapi/dir\"\n",
    "\n",
    "# Obtain the data loader.\n",
    "data_loader = get_loader(\n",
    "    transform=transform_train,\n",
    "    mode=\"train\",\n",
    "    batch_size=batch_size,\n",
    "    vocab_threshold=vocab_threshold,\n",
    "    vocab_from_file=False,\n",
    "    cocoapi_loc=cocoapi_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we run the code cell above, the data loader is stored in the variable `data_loader`.  \n",
    "\n",
    "We can now access the corresponding dataset as `data_loader.dataset`. This dataset is an instance of the `CoCoDataset` class in **coco_dataset.py**.  If you are unfamiliar with data loaders and datasets, you are encouraged to review [my post](http://www.sefidian.com/2022/03/09/writing-custom-datasets-and-dataloader-in-pytorch/) or [this PyTorch tutorial](http://pytorch.org/tutorials/beginner/data_loading_tutorial.html).\n",
    "\n",
    "### Exploring the `__getitem__` Method\n",
    "\n",
    "The `__getitem__` method in the `CoCoDataset` class determines how an image-caption pair is pre-processed before being incorporated into a batch.  This is true for all `Dataset` classes in PyTorch; if this is unfamiliar to you, please review [the tutorial linked above](http://pytorch.org/tutorials/beginner/data_loading_tutorial.html). \n",
    "\n",
    "When the data loader is in training mode, this method begins by first obtaining the filename (`path`) of a training image and its corresponding caption (`caption`).\n",
    "\n",
    "#### Image Pre-Processing \n",
    "\n",
    "Image pre-processing is relatively straightforward (from the `__getitem__` method in the `CoCoDataset` class):\n",
    "```python\n",
    "# Convert image to tensor and pre-process using transform\n",
    "image = Image.open(os.path.join(self.img_folder, path)).convert('RGB')\n",
    "image = self.transform(image)\n",
    "```\n",
    "After loading the image in the training folder with name `path`, the image is pre-processed using the same transform (`transform_train`) that was supplied when instantiating the data loader.  \n",
    "\n",
    "#### Caption Pre-Processing \n",
    "\n",
    "The captions also need to be pre-processed and prepped for training. In this example, for generating captions, we are aiming to create a model that predicts the next token of a sentence from previous tokens, so we turn the caption associated with any image into a list of tokenized words, before casting it to a PyTorch tensor that we can use to train the network.\n",
    "\n",
    "To understand in more detail how COCO captions are pre-processed, please take a look at the `vocab` instance variable of the `CoCoDataset` class. The code snippet below is pulled from the `__init__` method of the `CoCoDataset` class:\n",
    "```python\n",
    "def __init__(self, transform, mode, batch_size, vocab_threshold, vocab_file, start_word, \n",
    "        end_word, unk_word, annotations_file, vocab_from_file, img_folder):\n",
    "        ...\n",
    "        self.vocab = Vocabulary(vocab_threshold, vocab_file, start_word,\n",
    "            end_word, unk_word, annotations_file, vocab_from_file)\n",
    "        ...\n",
    "```\n",
    "From the code snippet above, we can see that `data_loader.dataset.vocab` is an instance of the `Vocabulary` class from **vocabulary.py**.  Take the time now to verify this by looking at the full code in **data_loader.py**.  \n",
    "\n",
    "We use this instance to pre-process the COCO captions (from the `__getitem__` method in the `CoCoDataset` class):\n",
    "\n",
    "```python\n",
    "# Convert caption to tensor of word ids.\n",
    "tokens = nltk.tokenize.word_tokenize(str(caption).lower())   # line 1\n",
    "caption = []                                                 # line 2\n",
    "caption.append(self.vocab(self.vocab.start_word))            # line 3\n",
    "caption.extend([self.vocab(token) for token in tokens])      # line 4\n",
    "caption.append(self.vocab(self.vocab.end_word))              # line 5\n",
    "caption = torch.Tensor(caption).long()                       # line 6\n",
    "```\n",
    "\n",
    "This code converts any string-valued caption to a list of integers, before casting it to a PyTorch tensor. To see how this code works, I'll apply it to the sample caption in the next code cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "sample_caption = \"A person doing a trick on a rail while riding a skateboard.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In **`line 1`** of the code snippet, every letter in the caption is converted to lowercase, and the [`nltk.tokenize.word_tokenize`](http://www.nltk.org/) function is used to obtain a list of string-valued tokens.  Run the next code cell to visualize the effect on `sample_caption`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['a', 'person', 'doing', 'a', 'trick', 'on', 'a', 'rail', 'while', 'riding', 'a', 'skateboard', '.']\n"
     ]
    }
   ],
   "source": [
    "sample_tokens = nltk.tokenize.word_tokenize(str(sample_caption).lower())\n",
    "print(sample_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**`line 2`** and **`line 3`** initializes an empty list and appends an integer to mark the start of a caption.  The [paper](https://arxiv.org/pdf/1411.4555.pdf) uses a special start word (and a special end word, which I'll examine below) to mark the beginning (and end) of a caption.\n",
    "\n",
    "This special start word (`\"<start>\"`) is decided when instantiating the data loader and is passed as a parameter (`start_word`). It is **required** to keep this parameter at its default value (`start_word=\"<start>\"`).\n",
    "\n",
    "As shown below, the integer `0` is always used to mark the start of a caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Special start word: <start>\n",
      "[0]\n"
     ]
    }
   ],
   "source": [
    "sample_caption = []\n",
    "\n",
    "start_word = data_loader.dataset.vocab.start_word\n",
    "print(\"Special start word:\", start_word)\n",
    "sample_caption.append(data_loader.dataset.vocab(start_word))\n",
    "print(sample_caption)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In **`line 4`**, we continue the list by adding integers that correspond to each of the tokens in the caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 3, 98, 754, 3, 396, 39, 3, 1010, 207, 139, 3, 753, 18]\n"
     ]
    }
   ],
   "source": [
    "sample_caption.extend([data_loader.dataset.vocab(token) for token in sample_tokens])\n",
    "print(sample_caption)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**`line 5`** appends a final integer to mark the end of the caption.  \n",
    "\n",
    "Identical to the case of the special start word (above), the special end word (`\"<end>\"`) is decided when instantiating the data loader and is passed as a parameter (`end_word`). It is **required** to keep this parameter at its default value (`end_word=\"<end>\"`).\n",
    "\n",
    "As shown below, the integer `1` is always used to  mark the end of a caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Special end word: <end>\n",
      "[0, 3, 98, 754, 3, 396, 39, 3, 1010, 207, 139, 3, 753, 18, 1]\n"
     ]
    }
   ],
   "source": [
    "end_word = data_loader.dataset.vocab.end_word\n",
    "print(\"Special end word:\", end_word)\n",
    "\n",
    "sample_caption.append(data_loader.dataset.vocab(end_word))\n",
    "print(sample_caption)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, in **`line 6`**, we convert the list of integers to a PyTorch tensor and cast it to [long type](http://pytorch.org/docs/master/tensors.html#torch.Tensor.long). Read more about the different types of PyTorch tensors on the [website](http://pytorch.org/docs/master/tensors.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([   0,    3,   98,  754,    3,  396,   39,    3, 1010,  207,  139,    3,\n",
      "         753,   18,    1])\n"
     ]
    }
   ],
   "source": [
    "sample_caption = torch.Tensor(sample_caption).long()\n",
    "print(sample_caption)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IDs\n",
    "start : 0\n",
    "\n",
    "end : 1\n",
    "\n",
    "unk : 2    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it!  In summary, any caption is converted to a list of tokens, with _special_ start and end tokens marking the beginning and end of the sentence:\n",
    "```\n",
    "[<start>, 'a', 'person', 'doing', 'a', 'trick', 'while', 'riding', 'a', 'skateboard', '.', <end>]\n",
    "```\n",
    "This list of tokens is then turned into a list of integers, where every distinct word in the vocabulary has an associated integer value:\n",
    "```\n",
    "[0, 3, 98, 754, 3, 396, 207, 139, 3, 753, 18, 1]\n",
    "```\n",
    "Finally, this list is converted to a PyTorch tensor.  All of the captions in the COCO dataset are pre-processed using this same procedure from **`lines 1-6`** described above.  \n",
    "\n",
    "As shown, in order to convert a token to its corresponding integer, we call `data_loader.dataset.vocab` as a function.  The details of how this call works can be explored in the `__call__` method in the `Vocabulary` class in **vocabulary.py**.  \n",
    "\n",
    "```python\n",
    "def __call__(self, word):\n",
    "    if not word in self.word2idx:\n",
    "        return self.word2idx[self.unk_word]\n",
    "    return self.word2idx[word]\n",
    "```\n",
    "\n",
    "The `word2idx` instance variable is a Python [dictionary](https://docs.python.org/3/tutorial/datastructures.html#dictionaries) that is indexed by string-valued keys (mostly tokens obtained from training captions). For each key, the corresponding value is the integer that the token is mapped to in the pre-processing step.\n",
    "\n",
    "Run cell below to view a subset of this dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'<start>': 0,\n",
       " '<end>': 1,\n",
       " '<unk>': 2,\n",
       " 'a': 3,\n",
       " 'very': 4,\n",
       " 'clean': 5,\n",
       " 'and': 6,\n",
       " 'well': 7,\n",
       " 'decorated': 8,\n",
       " 'empty': 9}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Preview the word2idx dictionary.\n",
    "dict(list(data_loader.dataset.vocab.word2idx.items())[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also print the total number of keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of tokens in vocabulary: 8852\n"
     ]
    }
   ],
   "source": [
    "# Print the total number of keys in the word2idx dictionary.\n",
    "print(\"Total number of tokens in vocabulary:\", len(data_loader.dataset.vocab))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In **vocabulary.py** codes, the `word2idx` dictionary is created by looping over the captions in the training dataset. If a token appears no less than `vocab_threshold` times in the training set, then it is added as a key to the dictionary and assigned a corresponding unique integer. We can amend the `vocab_threshold` argument when instantiating the data loader.  Note that in general, **smaller** values for `vocab_threshold` yield a **larger** number of tokens in the vocabulary. We can check this in the next code cell by decreasing the value of `vocab_threshold` before creating a new data loader.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Modify the minimum word count threshold.\n",
    "vocab_threshold = 4\n",
    "\n",
    "# Obtain the data loader.\n",
    "data_loader = get_loader(\n",
    "    transform=transform_train,\n",
    "    mode=\"train\",\n",
    "    batch_size=batch_size,\n",
    "    vocab_threshold=vocab_threshold,\n",
    "    vocab_from_file=False,\n",
    "    cocoapi_loc=cocoapi_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Print the total number of keys in the word2idx dictionary.\n",
    "print(f\"Total number of tokens in vocabulary: {len(data_loader.dataset.vocab)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are also a few special keys in the `word2idx` dictionary.  For example, the special start word (`\"<start>\"`) and special end word (`\"<end>\"`).  There is one more special token, corresponding to unknown words (`\"<unk>\"`).  All tokens that don't appear anywhere in the `word2idx` dictionary are considered unknown words. In the pre-processing step, any unknown tokens are mapped to the integer `2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Special unknown word: <unk>\n",
      "All unknown words are mapped to this integer: 2\n"
     ]
    }
   ],
   "source": [
    "unk_word = data_loader.dataset.vocab.unk_word\n",
    "print(f\"Special unknown word: {unk_word}\")\n",
    "\n",
    "print(\n",
    "    f\"All unknown words are mapped to this integer: {data_loader.dataset.vocab(unk_word)}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check this by pre-processing the provided nonsense words that never appear in the training captions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "print(data_loader.dataset.vocab(\"jfkafejw\"))\n",
    "print(data_loader.dataset.vocab(\"ieowoqjf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "18\n"
     ]
    }
   ],
   "source": [
    "print(data_loader.dataset.vocab(\".\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final thing to mention is the `vocab_from_file` argument that is supplied when creating a data loader.  To understand this argument, note that when a new data loader is created, the vocabulary (`data_loader.dataset.vocab`) is saved as a [pickle](https://docs.python.org/3/library/pickle.html) file in the project folder, with filename `vocab.pkl`.\n",
    "\n",
    "If you want to tweak the value of the `vocab_threshold` argument, you **must** set `vocab_from_file=False` to have the changes take effect.  \n",
    "\n",
    "Once the `vocab_threshold` argument has been chosen, run the data loader *one more time* with the chosen `vocab_threshold` to save the new vocabulary to file. Then, we can henceforth set `vocab_from_file=True` to load the vocabulary from file and speed the instantiation of the data loader. Note that building the vocabulary from scratch is the most time-consuming part of instantiating the data loader, and we set `vocab_from_file=True` as soon as we fix the configuration.\n",
    "\n",
    "Note that if `vocab_from_file=True`, then any supplied argument for `vocab_threshold` when instantiating the data loader is completely ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary successfully loaded from vocab.pkl file!\n",
      "loading annotations into memory...\n",
      "Done (t=0.56s)\n",
      "creating index...\n",
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 414113/414113 [00:43<00:00, 9587.01it/s]\n"
     ]
    }
   ],
   "source": [
    "# Obtain the data loader (from file). Note that it runs much faster than before!\n",
    "data_loader = get_loader(\n",
    "    transform=transform_train,\n",
    "    mode=\"train\",\n",
    "    batch_size=batch_size,\n",
    "    vocab_from_file=True,\n",
    "    cocoapi_loc=cocoapi_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The way to use the data loader to obtain batches of training data is explained in the next section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step2'></a>\n",
    "## Step 2: Using the Data Loader to Obtain Batches\n",
    "\n",
    "The captions in the dataset vary greatly in length. We can see this by examining `data_loader.dataset.caption_lengths`, a Python list with one entry for each training caption (where the value stores the length of the corresponding caption).  \n",
    "\n",
    "In the code cell below, we use this list to print the total number of captions in the training data with each length. As seen below, the majority of captions have length 10. Likewise, very short and very long captions are quite rare.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(list, 414113)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(data_loader.dataset.caption_lengths), len(data_loader.dataset.caption_lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "value: 10 --- count: 86302\n",
      "value: 11 --- count: 79971\n",
      "value:  9 --- count: 71920\n",
      "value: 12 --- count: 57653\n",
      "value: 13 --- count: 37668\n",
      "value: 14 --- count: 22342\n",
      "value:  8 --- count: 20742\n",
      "value: 15 --- count: 12839\n",
      "value: 16 --- count:  7736\n",
      "value: 17 --- count:  4845\n",
      "value: 18 --- count:  3101\n",
      "value: 19 --- count:  2017\n",
      "value:  7 --- count:  1594\n",
      "value: 20 --- count:  1453\n",
      "value: 21 --- count:   997\n",
      "value: 22 --- count:   684\n",
      "value: 23 --- count:   533\n",
      "value: 24 --- count:   384\n",
      "value: 25 --- count:   277\n",
      "value: 26 --- count:   214\n",
      "value: 27 --- count:   160\n",
      "value: 28 --- count:   114\n",
      "value: 29 --- count:    87\n",
      "value: 30 --- count:    58\n",
      "value: 31 --- count:    49\n",
      "value: 32 --- count:    44\n",
      "value: 34 --- count:    40\n",
      "value: 37 --- count:    32\n",
      "value: 35 --- count:    31\n",
      "value: 33 --- count:    30\n",
      "value: 36 --- count:    26\n",
      "value: 38 --- count:    18\n",
      "value: 39 --- count:    18\n",
      "value: 43 --- count:    16\n",
      "value: 44 --- count:    16\n",
      "value: 48 --- count:    12\n",
      "value: 45 --- count:    11\n",
      "value: 42 --- count:    10\n",
      "value: 40 --- count:     9\n",
      "value: 49 --- count:     9\n",
      "value: 46 --- count:     9\n",
      "value: 47 --- count:     7\n",
      "value: 50 --- count:     6\n",
      "value: 51 --- count:     6\n",
      "value: 41 --- count:     6\n",
      "value: 52 --- count:     5\n",
      "value: 54 --- count:     3\n",
      "value: 56 --- count:     2\n",
      "value:  6 --- count:     2\n",
      "value: 53 --- count:     2\n",
      "value: 55 --- count:     2\n",
      "value: 57 --- count:     1\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "# Tally the total number of training captions with each length.\n",
    "counter = Counter(data_loader.dataset.caption_lengths)\n",
    "lengths = sorted(counter.items(), key=lambda pair: pair[1], reverse=True)\n",
    "for value, count in lengths:\n",
    "    print(\"value: %2d --- count: %5d\" % (value, count))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate batches of training data, we begin by first sampling a caption length (where the probability that any length is drawn is proportional to the number of captions with that length in the dataset).  Then, we retrieve a batch of size `batch_size` of image-caption pairs, where all captions have the sampled length.  This approach for assembling batches matches the procedure in [this paper](https://arxiv.org/pdf/1502.03044.pdf) and has been shown to be computationally efficient without degrading performance.\n",
    "\n",
    "The code cell below generates a batch.  The `get_train_indices` method in the `CoCoDataset` class first samples a caption length, and then samples `batch_size` indices corresponding to training data points with captions of that length. These indices are stored below in `indices`.\n",
    "\n",
    "These indices are supplied to the data loader, which then is used to retrieve the corresponding data points. The pre-processed images and captions in the batch are stored in `images` and `captions`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "pycharm": {
     "is_executing": false
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n",
      "sampled indices: [388202, 286065, 397083, 408969, 92246, 354550, 13088, 43820, 124136, 374756]\n",
      "images.shape: torch.Size([10, 3, 224, 224])\n",
      "captions.shape: torch.Size([10, 13])\n"
     ]
    }
   ],
   "source": [
    "print(batch_size)\n",
    "# Randomly sample a caption length, and sample indices with that length.\n",
    "indices = data_loader.dataset.get_train_indices()\n",
    "print(\"sampled indices:\", indices)\n",
    "\n",
    "# Create and assign a batch sampler to retrieve a batch with the sampled indices.\n",
    "new_sampler = data.sampler.SubsetRandomSampler(indices=indices)\n",
    "data_loader.batch_sampler.sampler = new_sampler\n",
    "\n",
    "# Obtain the batch.\n",
    "images, captions = next(iter(data_loader))\n",
    "\n",
    "print(\"images.shape:\", images.shape)\n",
    "print(\"captions.shape:\", captions.shape)\n",
    "\n",
    "# Uncomment the lines of code below to print the pre-processed images and captions.\n",
    "# print('images:', images)\n",
    "# print('captions:', captions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Each time we run the code cell above, a different caption length is sampled**, and a different batch of training data is returned.\n",
    "\n",
    "We will train our model in the next notebook in this sequence (**2_Training.ipynb**).\n",
    "\n",
    "> Before moving to the next notebook in the sequence (**2_Training.ipynb**), take the time to become very familiar with the code in **coco_dataset.py**, **data_loader.py**, and **vocabulary.py**.  **Step 1** and **Step 2** of this notebook are designed to help facilitate a basic introduction and guide the understanding. However, our description is not exhaustive.\n",
    "\n",
    "In the next steps, we focus on learning how to specify a CNN-RNN architecture in PyTorch, towards the goal of image captioning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Architecture Details\n",
    "\n",
    "![Image Captioning CNN-RNN model](images/encoder-decoder.png)\n",
    "\n",
    "The architecture consists of a CNN encoder and RNN decoder. The CNN encoder is a pre-trained ResNet on ImageNet, which is a VGG convolutional neural network with skip connections. It has been proven to work really well on tasks like image recognition because the residual connections help model the residual differences before and after the convolution with the help of the identity block. A good pre-trained network on ImageNet is already good at extracting both useful low-level and high-level features for image tasks, so it naturally serves as a feature encoder for the image we want to caption. Since we are not doing the traditional image classification task, we drop the last fully connected layer and replace it without a new trainable fully connected layer to help transform the final feature map to an encoding that is more useful for the RNN decoder.\n",
    "\n",
    "RNNs have long been shown useful in language tasks due to their ability to model data with sequential nature, such as language. Specifically, LSTMs even incorporate both long-term and short-term information as memories in the network. Thus, we pick an RNN decoder for the captioning task. Specifically, following the spirit of sequence to sequence (seq2seq) models used in translation, I leveraged the architecture choices in [this paper](https://arxiv.org/pdf/1411.4555.pdf) to use an LSTM to generate captions based on the encoded information from the CNN encoder. Specifically, **I first use the CNN encoder output concatenated with the \"START\" token as the initial input for the RNN decoder.** I apply a fully connected layer on the hidden states at that timestamp to output a softmax probability over the words in our entire vocabulary, where we choose the word with the highest probability as the word generated at that timestamp. Then, we feed this predicted word back again as the input for the next step. We continue so until we generated a caption of max length, or the network generated the \"STOP\" token, which indicates the end of the sentence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step3'></a>\n",
    "## Step 3: Experimenting with the CNN Encoder\n",
    "\n",
    "Run the code cell below to import `EncoderCNN` and `DecoderRNN` from **model.py**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Import EncoderCNN and DecoderRNN.\n",
    "# Watch for any changes in model.py, and re-load it automatically.\n",
    "from model import EncoderCNN, DecoderRNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next code cell we define a `device` that we will use move PyTorch tensors to GPU (if CUDA is available).  Run this code cell before continuing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the code cell below to instantiate the CNN encoder in `encoder`.  \n",
    "\n",
    "The pre-processed images from the batch in **Step 2** of this notebook are then passed through the encoder, and the output is stored in `features`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type(features): <class 'torch.Tensor'>\n",
      "features.shape: torch.Size([10, 256])\n",
      "captions.shape: torch.Size([10, 13])\n"
     ]
    }
   ],
   "source": [
    "# Specify the dimensionality of the image embedding.\n",
    "image_embed_size = 256\n",
    "\n",
    "# Initialize the encoder.\n",
    "encoder = EncoderCNN(image_embed_size)\n",
    "\n",
    "# Move the encoder to appropriate device.\n",
    "encoder.to(device)\n",
    "\n",
    "# Move last batch of images (from Step 2) to GPU if CUDA is available.\n",
    "images = images.to(device)\n",
    "\n",
    "# Pass the images through the encoder.\n",
    "features = encoder(images)\n",
    "\n",
    "print(\"type(features):\", type(features))\n",
    "print(\"features.shape:\", features.shape)\n",
    "print(\"captions.shape:\", captions.shape)\n",
    "\n",
    "# Check that the encoder satisfies the requirements!\n",
    "assert type(features) == torch.Tensor, \"Encoder output needs to be a PyTorch Tensor.\"\n",
    "\n",
    "assert (features.shape[0] == batch_size) and (\n",
    "    features.shape[1] == image_embed_size\n",
    "), \"The shape of the encoder output is incorrect.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The encoder uses the pre-trained ResNet-50 architecture (with the final fully-connected layer removed) to extract features from a batch of pre-processed images. The output is then flattened to a vector, before being passed through a `Linear` layer to transform the feature vector to have the same size as the word embedding.\n",
    "\n",
    "![Encoder](images/encoder.png)\n",
    "\n",
    "We can amend the encoder in **model.py**, to experiment with other architectures. In particular, using a [different pre-trained model architecture](http://pytorch.org/docs/master/torchvision/models.html) and adding [add batch normalization](http://pytorch.org/docs/master/nn.html#normalization-layers) could be good options.  \n",
    "\n",
    "\n",
    "For this project, I incorporated a pre-trained CNN into the encoder. The `EncoderCNN` class takes `image_embed_size` as an input argument, which will also correspond to the dimensionality of the input to the RNN decoder that we will implement in Step 4. When we train the model in the next notebook in this sequence (**2_Training.ipynb**), we can tweak the value of `image_embed_size`.\n",
    "\n",
    "If you decide to modify the `EncoderCNN` class, save **model.py** and re-execute the code cell above. If the code cell returns an assertion error, then please follow the instructions to modify the code before proceeding. The assert statements ensure that `features` is a PyTorch tensor with shape `[batch_size, image_embed_size]`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step4'></a>\n",
    "## Step 4: Implementing the RNN Decoder\n",
    "\n",
    "Before executing the next code cell, please read the `__init__` and `forward` methods in the `DecoderRNN` class in **model.py**. We will work on `sample` method when we reach **3_Inference.ipynb**.\n",
    "\n",
    "The decoder is an instance of the `DecoderRNN` class and accepts the followings as input:\n",
    "- the PyTorch tensor `features` containing the embedded image features (outputted in Step 3, when the last batch of images from Step 2 was passed through `encoder`)\n",
    "- a PyTorch tensor corresponding to the last batch of captions (`captions`) from Step 2.\n",
    "\n",
    "Note that the way I have written the data loader simplifies the code a bit. In particular, every training batch will contain pre-processed **captions where all have the same length** (`captions.shape[1]`), so **we do not need to worry about padding**.  \n",
    "> I have implemented the decoder described in [this paper](https://arxiv.org/pdf/1411.4555.pdf), you can implement any architecture of your choosing.  \n",
    "\n",
    "Although we will test the decoder using the last batch that is currently stored in the notebook, the decoder can accept an arbitrary batch (of embedded image features and pre-processed captions [where all captions have the same length]) as input.  \n",
    "\n",
    "![Decoder](images/decoder.png)\n",
    "\n",
    "In the code cell below, `outputs` is a PyTorch tensor with size `[batch_size, captions.shape[1], vocab_size]`.  The output designed such that `outputs[i,j,k]` contains the model's predicted score, indicating how likely the `j`-th token in the `i`-th caption in the batch is the `k`-th token in the vocabulary. In the next notebook of the sequence (**2_Training.ipynb**), we provide code to supply these scores to the [`torch.nn.CrossEntropyLoss`](http://pytorch.org/docs/master/nn.html#torch.nn.CrossEntropyLoss) optimizer in PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "256\n"
     ]
    }
   ],
   "source": [
    "print(image_embed_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type(outputs): <class 'torch.Tensor'>\n",
      "outputs.shape: torch.Size([10, 13, 8852])\n"
     ]
    }
   ],
   "source": [
    "# Specify the number of features in the hidden state of the RNN decoder.\n",
    "hidden_size = 512\n",
    "\n",
    "word_embed_size = image_embed_size\n",
    "\n",
    "# Store the size of the vocabulary.\n",
    "vocab_size = len(data_loader.dataset.vocab)\n",
    "\n",
    "# Initialize the decoder.\n",
    "decoder = DecoderRNN(word_embed_size, hidden_size, vocab_size)\n",
    "\n",
    "# Move the decoder to proper device.\n",
    "decoder.to(device)\n",
    "\n",
    "\n",
    "# Move last batch of captions (from Step 1) to GPU if CUDA is available\n",
    "captions = captions.to(device)\n",
    "\n",
    "# Pass the encoder output and captions through the decoder.\n",
    "# outputs[i,j,k] contains the model's predicted score:\n",
    "# how likely the j-th token in the i-th caption in the batch is the k-th token in the vocabulary.\n",
    "\n",
    "outputs = decoder(features, captions)  # (bs, cap_length, vocab_size)\n",
    "\n",
    "\n",
    "print(\"type(outputs):\", type(outputs))\n",
    "print(\"outputs.shape:\", outputs.shape)\n",
    "\n",
    "# Check that the decoder satisfies the requirements!\n",
    "assert type(outputs) == torch.Tensor, \"Decoder output needs to be a PyTorch Tensor.\"\n",
    "assert (\n",
    "    (outputs.shape[0] == batch_size)\n",
    "    and (outputs.shape[1] == captions.shape[1])\n",
    "    and (outputs.shape[2] == vocab_size)\n",
    "), \"The shape of the decoder output is incorrect.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we are training the model in the next notebook in this sequence (**2_Training.ipynb**), we can tweak the value of `hidden_size`."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.7.10"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}