1_Preliminaries.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Captioning\n",
    "\n",
    "## Part 1: Load and Pre-Process Data and Experiment with Models\n",
    "\n",
    "---\n",
    "\n",
    "In this notebook, we will learn how to load and pre-process data from the [COCO dataset](http://cocodataset.org/#home). We will also experiment with a CNN-RNN model for automatically generating image captions. These are *not* the final models that we will use. For the final ones, see **model.py**.\n",
    "\n",
    "Use the links below to navigate the 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: Explore the Data Loader\n",
    "\n",
    "We will use a [data loader](http://pytorch.org/docs/master/data.html#torch.utils.data.DataLoader) to load the COCO dataset in batches. \n",
    "\n",
    "In the code cell below, we will initialize the data loader by using the `get_loader` function in **data_loader.py**.  \n",
    "\n",
    "The `get_loader` function takes as input a number of arguments that can be explored in **data_loader.py**. Most of the arguments must be left at their default values; we may amend the values of the arguments below:\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.\n",
    "2. **`mode`** - one of `'train'`, `'val'` (loads the training or validation data in batches) or `'test'` (for the test data). We will say that the data loader is in training, validation or test mode, respectively.\n",
    "3. **`batch_size`** - determines the batch size.  When training/validating the model, this is number of image-caption pairs used to amend the model weights in each training/validation step.\n",
    "4. **`vocab_threshold`** - the total number of times that a word must appear 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Watch for any changes in vocabulary.py, data_loader.py or model.py, and re-load it automatically.\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.64s)\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.60s)\n",
      "creating index...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1577/414113 [00:00<00:52, 7882.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 414113/414113 [00:48<00:00, 8624.65it/s]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from data_loader import get_loader\n",
    "from torchvision import transforms\n",
    "\n",
    "# Define a transform to pre-process the training images.\n",
    "transform_train = transforms.Compose([ \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((0.485, 0.456, 0.406),      # normalize image for pre-trained model\n",
    "                         (0.229, 0.224, 0.225))])\n",
    "\n",
    "# Set the minimum word count threshold.\n",
    "vocab_threshold = 5\n",
    "\n",
    "# Specify the batch size.\n",
    "batch_size = 10\n",
    "\n",
    "# Obtain the data loader.\n",
    "data_loader = get_loader(transform=transform_train,\n",
    "                         mode='train',\n",
    "                         batch_size=batch_size,\n",
    "                         vocab_threshold=vocab_threshold,\n",
    "                         vocab_from_file=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we ran the code cell above, the data loader was stored in the variable `data_loader`.  \n",
    "\n",
    "We can access the corresponding dataset as `data_loader.dataset`.  This dataset is an instance of the `CoCoDataset` class in **data_loader.py**.  For more information on data loaders and datasets see [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.  When the data loader is in training or validation mode, this method begins by first obtaining the filename (`path`) of an 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 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, we'll first need to 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",
    "`data_loader.dataset.vocab` is an instance of the `Vocabulary` class from **vocabulary.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",
    "As we will see soon, 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, we'll apply it to the sample caption in the next code cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['a', 'person', 'doing', 'a', 'trick', 'on', 'a', 'rail', 'while', 'riding', 'a', 'skateboard', '.']\n"
     ]
    }
   ],
   "source": [
    "import nltk\n",
    "\n",
    "sample_tokens = nltk.tokenize.word_tokenize(str(sample_caption).lower())\n",
    "print(sample_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In **`line 2`** and **`line 3`** we initialize an empty list and append an integer to mark the start of a caption.  This [paper](https://arxiv.org/pdf/1411.4555.pdf) uses a special start word (and a special end word, which we'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`).  We will keep this parameter at its default value (`start_word=\"<start>\"`).\n",
    "\n",
    "As we will see below, the integer `0` is always used to mark the start of a caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 3, 98, 754, 3, 396, 39, 3, 1009, 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": [
    "In **`line 5`**, we append 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`).  We keep this parameter at its default value (`end_word=\"<end>\"`).\n",
    "\n",
    "As we will see below, the integer `1` is always used to  mark the end of a caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Special end word: <end>\n",
      "[0, 3, 98, 754, 3, 396, 39, 3, 1009, 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).  More information about the different types of PyTorch tensors is available on the [website](http://pytorch.org/docs/master/tensors.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([    0,     3,    98,   754,     3,   396,    39,     3,  1009,\n",
      "          207,   139,     3,   753,    18,     1])\n"
     ]
    }
   ],
   "source": [
    "sample_caption = torch.Tensor(sample_caption).long()\n",
    "print(sample_caption)"
   ]
  },
  {
   "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 we saw, 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 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",
    "Use the code cell below to view a subset of this dictionary.  We also print the total number of keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'<start>': 0, '<end>': 1, '<unk>': 2, 'a': 3, 'very': 4, 'clean': 5, 'and': 6, 'well': 7, 'decorated': 8, 'empty': 9}\n",
      "Total number of tokens in vocabulary: 8855\n"
     ]
    }
   ],
   "source": [
    "# Preview the word2idx dictionary.\n",
    "print (dict(list(data_loader.dataset.vocab.word2idx.items())[:10]))\n",
    "\n",
    "# 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**, 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.  In general, **smaller** values for `vocab_threshold` yield a **larger** number of tokens in the vocabulary.  We can see this in the next two code cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.70s)\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.58s)\n",
      "creating index...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 531/414113 [00:00<01:17, 5306.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 414113/414113 [00:49<00:00, 8370.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of tokens in vocabulary: 8855\n"
     ]
    }
   ],
   "source": [
    "# Minimum word count threshold.\n",
    "vocab_threshold = 5\n",
    "\n",
    "# Obtain the data loader.\n",
    "data_loader = get_loader(transform=transform_train,\n",
    "                         mode='train',\n",
    "                         batch_size=batch_size,\n",
    "                         vocab_threshold=vocab_threshold,\n",
    "                         vocab_from_file=False)\n",
    "# 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": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.58s)\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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 28/414113 [00:00<35:17, 195.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done (t=0.55s)\n",
      "creating index...\n",
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 414113/414113 [00:49<00:00, 8372.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of tokens in vocabulary: 6335\n"
     ]
    }
   ],
   "source": [
    "# Minimum word count threshold.\n",
    "vocab_threshold = 10\n",
    "\n",
    "# Obtain the data loader.\n",
    "data_loader = get_loader(transform=transform_train,\n",
    "                         mode='train',\n",
    "                         batch_size=batch_size,\n",
    "                         vocab_threshold=vocab_threshold,\n",
    "                         vocab_from_file=False)\n",
    "\n",
    "# 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": [
    "There are also a few special keys in the `word2idx` dictionary.  Other than 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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Special unknown word: <unk>\n",
      "All unknown words are mapped to this integer: 2\n",
      "For example:\n",
      "'jfkafejw' is mapped to 2\n"
     ]
    }
   ],
   "source": [
    "unk_word = data_loader.dataset.vocab.unk_word\n",
    "print('Special unknown word:', unk_word)\n",
    "\n",
    "print('All unknown words are mapped to this integer:', data_loader.dataset.vocab(unk_word))\n",
    "print (\"For example:\")\n",
    "print(\"'jfkafejw' is mapped to\", data_loader.dataset.vocab('jfkafejw'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final thing to mention is the `vocab_from_file` argument that is supplied when creating a data loader.  When we create a new data loader, 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 we are still tweaking the value of the `vocab_threshold` argument, we **must** set `vocab_from_file=False` to have our changes take effect.  \n",
    "\n",
    "But once we are happy with the value that we have chosen for the `vocab_threshold` argument, we need only run the data loader *one more time* with our 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 so we are strongly encouraged to set `vocab_from_file=True` as soon as we are able.\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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary successfully loaded from vocab.pkl file!\n",
      "loading annotations into memory...\n",
      "Done (t=0.71s)\n",
      "creating index...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 516/414113 [00:00<01:20, 5156.22it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 414113/414113 [00:50<00:00, 8164.98it/s]\n"
     ]
    }
   ],
   "source": [
    "# Obtain the data loader (from file). Note that it runs much faster than before!\n",
    "data_loader = get_loader(transform=transform_train,\n",
    "                         mode='train',\n",
    "                         batch_size=batch_size,\n",
    "                         vocab_from_file=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step2'></a>\n",
    "## Step 2: Use 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 we will see 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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "value: 10 --- count: 86334\n",
      "value: 11 --- count: 79948\n",
      "value:  9 --- count: 71934\n",
      "value: 12 --- count: 57637\n",
      "value: 13 --- count: 37645\n",
      "value: 14 --- count: 22335\n",
      "value:  8 --- count: 20771\n",
      "value: 15 --- count: 12841\n",
      "value: 16 --- count:  7729\n",
      "value: 17 --- count:  4842\n",
      "value: 18 --- count:  3104\n",
      "value: 19 --- count:  2014\n",
      "value:  7 --- count:  1597\n",
      "value: 20 --- count:  1451\n",
      "value: 21 --- count:   999\n",
      "value: 22 --- count:   683\n",
      "value: 23 --- count:   534\n",
      "value: 24 --- count:   383\n",
      "value: 25 --- count:   277\n",
      "value: 26 --- count:   215\n",
      "value: 27 --- count:   159\n",
      "value: 28 --- count:   115\n",
      "value: 29 --- count:    86\n",
      "value: 30 --- count:    58\n",
      "value: 31 --- count:    49\n",
      "value: 32 --- count:    44\n",
      "value: 34 --- count:    39\n",
      "value: 37 --- count:    32\n",
      "value: 33 --- count:    31\n",
      "value: 35 --- count:    31\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",
    "Run the code cell below to generate a batch.  The `get_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": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10 sampled indices: [258091, 379547, 124568, 400930, 16735, 279034, 111811, 24670, 399565, 247309]\n",
      "images.shape: torch.Size([10, 3, 224, 224])\n",
      "captions.shape: torch.Size([10, 13])\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch.utils.data as data\n",
    "\n",
    "# Randomly sample a caption length, and sample indices with that length.\n",
    "indices = data_loader.dataset.get_indices()\n",
    "print('{} sampled indices: {}'.format(len(indices), indices))\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",
    "for batch in data_loader:\n",
    "    images, captions = batch[0], batch[1]\n",
    "    break\n",
    "    \n",
    "print('images.shape:', images.shape)\n",
    "print('captions.shape:', captions.shape)\n",
    "\n",
    "# Print the pre-processed images and captions.\n",
    "#print('images:', images)\n",
    "#print('captions:', captions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step3'></a>\n",
    "## Step 3: Experiment with the CNN Encoder\n",
    "\n",
    "First, we will import `EncoderCNN` and `DecoderRNN` from **model.py**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import EncoderCNN and DecoderRNN. \n",
    "from model import EncoderCNN, DecoderRNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will 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`. The assert statement ensures that `features` has shape `[batch_size, embed_size]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type(features): <class 'torch.Tensor'>\n",
      "features.shape: torch.Size([10, 256])\n"
     ]
    }
   ],
   "source": [
    "# Specify the dimensionality of the image embedding.\n",
    "embed_size = 256\n",
    "\n",
    "# Initialize the encoder. (We can add additional arguments if necessary.)\n",
    "encoder = EncoderCNN(embed_size)\n",
    "\n",
    "# Move the encoder to GPU if CUDA is available.\n",
    "if torch.cuda.is_available():\n",
    "    encoder = encoder.cuda()\n",
    "    \n",
    "# Move the last batch of images from Step 2 to GPU if CUDA is available\n",
    "if torch.cuda.is_available():\n",
    "    images = images.cuda()\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",
    "\n",
    "# Check that our encoder satisfies some requirements of the project!\n",
    "assert (features.shape[0]==batch_size) & (features.shape[1]==embed_size), \"The shape of the encoder output is incorrect.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This 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 could amend the encoder in **model.py**, to experiment with other architectures, such as using a [different pre-trained model architecture](http://pytorch.org/docs/master/torchvision/models.html) or [adding batch normalization](http://pytorch.org/docs/master/nn.html#normalization-layers).  \n",
    "\n",
    "For this project, we will **incorporate a pre-trained CNN into our encoder**.  The `EncoderCNN` class must take `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 our model in the next notebook in this sequence (**2_Training.ipynb**), we will tweak the value of `embed_size`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step4'></a>\n",
    "## Step 4: Implement the RNN Decoder\n",
    "\n",
    "Our decoder will be an instance of the `DecoderRNN` class from **model.py** and must accept 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`), along with\n",
    "- a PyTorch tensor corresponding to the last batch of captions (`captions`) from Step 2.\n",
    "\n",
    "Every training batch will contain pre-processed captions where all have the same length (`captions.shape[1]`), so **we won't need to worry about padding**.  \n",
    "\n",
    "Although we will test the decoder using the last batch that is currently stored in the notebook, our decoder should 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` should have size `[batch_size, captions.shape[1], vocab_size]`.  Our output should be 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 will 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": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type(outputs): <class 'torch.Tensor'>\n",
      "outputs.shape: torch.Size([10, 13, 6335])\n"
     ]
    }
   ],
   "source": [
    "# Specify the number of features in the hidden state of the RNN decoder.\n",
    "hidden_size = 512\n",
    "\n",
    "# Store the size of the vocabulary.\n",
    "vocab_size = len(data_loader.dataset.vocab)\n",
    "\n",
    "# Initialize the decoder.\n",
    "decoder = DecoderRNN(embed_size, hidden_size, vocab_size)\n",
    "# Move the decoder to GPU if CUDA is available.\n",
    "if torch.cuda.is_available():\n",
    "    decoder = decoder.cuda()\n",
    "    \n",
    "# Move the last batch of captions (from Step 1) to GPU if cuda is availble \n",
    "if torch.cuda.is_available():\n",
    "    captions = captions.cuda()\n",
    "# Pass the encoder output and captions through the decoder\n",
    "outputs = decoder(features, captions)\n",
    "\n",
    "print('type(outputs):', type(outputs))\n",
    "print('outputs.shape:', outputs.shape)\n",
    "\n",
    "# Check that our decoder satisfies some requirements of the project!\n",
    "assert (outputs.shape[0]==batch_size) & (outputs.shape[1]==captions.shape[1]) & (outputs.shape[2]==vocab_size), \"The shape of the decoder output is incorrect.\""
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}