1_Preliminaries.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computer Vision Nanodegree\n",
    "\n",
    "## Project: Image Captioning\n",
    "\n",
    "---\n",
    "\n",
    "In this notebook, you will learn how to load and pre-process data from the [COCO dataset](http://cocodataset.org/#home). You will also design a CNN-RNN model for automatically generating image captions.\n",
    "\n",
    "Note that **any amendments that you make to this notebook will not be graded**.  However, you will use the instructions provided in **Step 3** and **Step 4** to implement your own CNN encoder and RNN decoder by making amendments to the **models.py** file provided as part of this project.  Your **models.py** file **will be graded**. \n",
    "\n",
    "Feel free to 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 have already written a [data loader](http://pytorch.org/docs/master/data.html#torch.utils.data.DataLoader) that you can use to load the COCO dataset in batches. \n",
    "\n",
    "In the code cell below, you will initialize the data loader by using the `get_loader` function in **data_loader.py**.  \n",
    "\n",
    "> For this project, you are not permitted to change the **data_loader.py** file, which must be used as-is.\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** in a new window.  Most of the arguments must be left at their default values, and you are only allowed to 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.  For now, you are encouraged to keep the transform as provided in `transform_train`.  You will have the opportunity later to choose your own image transform to pre-process the COCO images.\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.  While following the instructions in this notebook, please keep the data loader in training mode 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 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.  Be patient - it may take a couple of minutes to run!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/opt/cocoapi/PythonAPI')\n",
    "from pycocotools.coco import COCO\n",
    "!pip install nltk\n",
    "import nltk\n",
    "nltk.download('punkt')\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 you ran the code cell above, the data loader was stored in the variable `data_loader`.  \n",
    "\n",
    "You can access the corresponding dataset as `data_loader.dataset`.  This dataset is an instance of the `CoCoDataset` class in **data_loader.py**.  If you are unfamiliar with data loaders and datasets, you are encouraged to review [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, 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",
    "From the code snippet above, you can see that `data_loader.dataset.vocab` is an instance of the `Vocabulary` class from **vocabulary.py**.  Take the time now to verify this for yourself 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",
    "As you 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": null,
   "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.  Run the next code cell to visualize the effect on `sample_caption`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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.  The [paper](https://arxiv.org/pdf/1411.4555.pdf) that you are encouraged to implement 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`).  You are **required** to keep this parameter at its default value (`start_word=\"<start>\"`).\n",
    "\n",
    "As you will see below, the integer `0` is always used to mark the start of a caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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`).  You are **required** to keep this parameter at its default value (`end_word=\"<end>\"`).\n",
    "\n",
    "As you will see below, the integer `1` is always used to  mark the end of a caption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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).  You can read more about the different types of PyTorch tensors on the [website](http://pytorch.org/docs/master/tensors.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "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 you 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](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",
    "Use the code cell below to view a subset of this dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": [
    "As you will see if you examine the code 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.  You will have the option later to amend the `vocab_threshold` argument when instantiating your data loader.  Note that in general, **smaller** values for `vocab_threshold` yield a **larger** number of tokens in the vocabulary.  You are encouraged to check this for yourself 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": {},
   "outputs": [],
   "source": [
    "# Modify the minimum word count threshold.\n",
    "vocab_threshold = 4\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": [
    "There are also a few special keys in the `word2idx` dictionary.  You are already familiar with 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": null,
   "metadata": {},
   "outputs": [],
   "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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check this for yourself below, by pre-processing the provided nonsense words that never appear in the training captions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(data_loader.dataset.vocab('jfkafejw'))\n",
    "print(data_loader.dataset.vocab('ieowoqjf'))"
   ]
  },
  {
   "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 you 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 you are still tweaking the value of the `vocab_threshold` argument, you **must** set `vocab_from_file=False` to have your changes take effect.  \n",
    "\n",
    "But once you are happy with the value that you have chosen for the `vocab_threshold` argument, you need only run the data loader *one more time* with your chosen `vocab_threshold` to save the new vocabulary to file.  Then, you 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 you are strongly encouraged to set `vocab_from_file=True` as soon as you 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": null,
   "metadata": {},
   "outputs": [],
   "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": [
    "In the next section, you will learn how to use the data loader to obtain batches of training data."
   ]
  },
  {
   "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.  You 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 you 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": null,
   "metadata": {},
   "outputs": [],
   "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_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": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "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_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",
    "# (Optional) 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 you run the code cell above, a different caption length is sampled, and a different batch of training data is returned.  Run the code cell multiple times to check this out!\n",
    "\n",
    "You will train your model in the next notebook in this sequence (**2_Training.ipynb**). This code for generating training batches will be provided to you.\n",
    "\n",
    "> Before moving to the next notebook in the sequence (**2_Training.ipynb**), you are strongly encouraged to take the time to become very familiar with the code in  **data_loader.py** and **vocabulary.py**.  **Step 1** and **Step 2** of this notebook are designed to help facilitate a basic introduction and guide your understanding.  However, our description is not exhaustive, and it is up to you (as part of the project) to learn how to best utilize these files to complete the project.  __You should NOT amend any of the code in either *data_loader.py* or *vocabulary.py*.__\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": [
    "<a id='step3'></a>\n",
    "## Step 3: Experiment with the CNN Encoder\n",
    "\n",
    "Run the code cell below to import `EncoderCNN` and `DecoderRNN` from **model.py**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Watch for any changes in model.py, and re-load it automatically.\n",
    "% load_ext autoreload\n",
    "% autoreload 2\n",
    "\n",
    "# Import EncoderCNN and DecoderRNN. \n",
    "from model import EncoderCNN, DecoderRNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next code cell we define a `device` that you will use move PyTorch tensors to GPU (if CUDA is available).  Run this code cell before continuing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the dimensionality of the image embedding.\n",
    "embed_size = 256\n",
    "\n",
    "#-#-#-# Do NOT modify the code below this line. #-#-#-#\n",
    "\n",
    "# Initialize the encoder. (Optional: Add additional arguments if necessary.)\n",
    "encoder = EncoderCNN(embed_size)\n",
    "\n",
    "# Move the encoder to GPU if CUDA is available.\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",
    "\n",
    "# Check that your encoder satisfies some requirements of the project! :D\n",
    "assert type(features)==torch.Tensor, \"Encoder output needs to be a PyTorch Tensor.\" \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": [
    "The encoder that we provide to you 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",
    "You are welcome (and encouraged) to amend the encoder in **model.py**, to experiment with other architectures.  In particular, consider using a [different pre-trained model architecture](http://pytorch.org/docs/master/torchvision/models.html).  You may also like to [add batch normalization](http://pytorch.org/docs/master/nn.html#normalization-layers).  \n",
    "\n",
    "> You are **not** required to change anything about the encoder.\n",
    "\n",
    "For this project, you **must** incorporate a pre-trained CNN into your encoder.  Your `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 you will implement in Step 4.  When you train your model in the next notebook in this sequence (**2_Training.ipynb**), you are welcome to tweak the value of `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 your code before proceeding.  The assert statements ensure that `features` is a PyTorch tensor with shape `[batch_size, embed_size]`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='step4'></a>\n",
    "## Step 4: Implement the RNN Decoder\n",
    "\n",
    "Before executing the next code cell, you must write `__init__` and `forward` methods in the `DecoderRNN` class in **model.py**.  (Do **not** write the `sample` method yet - you will work with this method when you reach **3_Inference.ipynb**.)\n",
    "\n",
    "> The `__init__` and `forward` methods in the `DecoderRNN` class are the only things that you **need** to modify as part of this notebook.  You will write more implementations in the notebooks that appear later in the sequence.\n",
    "\n",
    "Your decoder will be an instance of the `DecoderRNN` class 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",
    "Note that the way we have written the data loader should simplify your code a bit.  In particular, every training batch will contain pre-processed captions where all have the same length (`captions.shape[1]`), so **you do not need to worry about padding**.  \n",
    "> While you are encouraged to implement the decoder described in [this paper](https://arxiv.org/pdf/1411.4555.pdf), you are welcome to implement any architecture of your choosing, as long as it uses at least one RNN layer, with hidden dimension `hidden_size`.  \n",
    "\n",
    "Although you will test the decoder using the last batch that is currently stored in the notebook, your decoder should be written to 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 be a PyTorch tensor with size `[batch_size, captions.shape[1], vocab_size]`.  Your 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 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the number of features in the hidden state of the RNN decoder.\n",
    "hidden_size = 512\n",
    "\n",
    "#-#-#-# Do NOT modify the code below this line. #-#-#-#\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",
    "\n",
    "# Move the decoder to GPU if CUDA is available.\n",
    "decoder.to(device)\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 = decoder(features, captions)\n",
    "\n",
    "print('type(outputs):', type(outputs))\n",
    "print('outputs.shape:', outputs.shape)\n",
    "\n",
    "# Check that your decoder satisfies some requirements of the project! :D\n",
    "assert type(outputs)==torch.Tensor, \"Decoder output needs to be a PyTorch Tensor.\"\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.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you train your model in the next notebook in this sequence (**2_Training.ipynb**), you are welcome to tweak the value of `hidden_size`."
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}