{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recurrent Neural Networks\n",
    "\n",
    "In this lab we will experiment with recurrent neural networks. We will build a text generation model that predicts a word given the previous words, and hence will allow us to generate a sentence. This can easily be extended to generating a sentence description for a given input image. RNNs are a useful type of model for predicting sequences or handling sequences of things as inputs. In this lab we will use again Pytorch's nn library.\n",
    "\n",
    "We will also be using the COCO dataset which includes images + textual descriptions (captions) + other annotations. We can browse the dataset here: http://cocodataset.org/#home\n",
    "\n",
    "First, let's import libraries and make sure we have everything properly installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch, json, string\n",
    "import torch.nn as nn\n",
    "import torch.autograd\n",
    "from torch.autograd import Variable\n",
    "import torch.utils.data as data\n",
    "import numpy as np\n",
    "from nltk.tokenize import word_tokenize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Loading and Preprocessing the Text\n",
    "Pytorch comes with a Dataset class for the COCO dataset but I will write my own class here. This class does two important things: 1) Building a vocabulary with the most frequent words, 2) Building utilities to convert a sentence into a list of word ids, and back. We are not going to be using the images for the purposes of the lab but you will use them in the assignment questions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "class CocoCaptions(data.Dataset):\n",
    "    # Load annotations in the initialization of the object.\n",
    "    def __init__(self, captionsFile, vocabulary = None):\n",
    "        self.data = json.load(open(captionsFile))\n",
    "        self.imageIds = self.data['images']\n",
    "        self.annotations = self.data['annotations']\n",
    "        \n",
    "        # Build a vocabulary if not provided.\n",
    "        if not vocabulary:\n",
    "            self.build_vocabulary()\n",
    "        else:\n",
    "            self.vocabulary = vocabulary\n",
    "        \n",
    "    # Build a vocabulary using the top 5000 words.\n",
    "    def build_vocabulary(self, vocabularySize = 5000):\n",
    "        # Count words, this will take a while.\n",
    "        word_counter = dict()\n",
    "        for annotation in tqdm(self.annotations, desc = 'Building vocabulary'):\n",
    "            words = word_tokenize(annotation['caption'].lower())\n",
    "            for word in words:\n",
    "                word_counter[word] = word_counter.get(word, 0) + 1\n",
    "                \n",
    "        # Sort the words and find keep only the most frequent words.\n",
    "        sorted_words = sorted(list(word_counter.items()), \n",
    "                              key = lambda x: -x[1])\n",
    "        most_frequent_words = [w for (w, c) in sorted_words[:vocabularySize]]\n",
    "        word2id = {w: (index + 1) for (index, w) in enumerate(most_frequent_words)}\n",
    "        \n",
    "        # Add a special characters for START, END sentence, and UNKnown words.\n",
    "        word2id['[END]'] = 0\n",
    "        word2id['[START]'] = len(word2id)\n",
    "        word2id['UNK'] = len(word2id)\n",
    "        id2word = {index: w for (w, index) in word2id.items()}\n",
    "        self.vocabulary = {'word2id': word2id, 'id2word': id2word}\n",
    "    \n",
    "    # Transform a caption into a list of word ids.\n",
    "    def caption2ids(self, caption):\n",
    "        word2id = self.vocabulary['word2id']\n",
    "        caption_ids = [word2id.get(w, word2id['UNK']) for w in word_tokenize(caption.lower())]\n",
    "        caption_ids.insert(0, word2id['[START]'])\n",
    "        caption_ids.append(word2id['[END]'])\n",
    "        return torch.LongTensor(caption_ids)\n",
    "    \n",
    "    # Transform a list of word ids into a caption.\n",
    "    def ids2caption(self, caption_ids):\n",
    "        id2word = self.vocabulary['id2word']\n",
    "        return string.join([id2word[w] for w in caption_ids], \" \")\n",
    "    \n",
    "    # Return imgId, and a random caption for that image.\n",
    "    def __getitem__(self, index):\n",
    "        annotation = self.annotations[index]\n",
    "        return annotation['image_id'], self.caption2ids(annotation['caption'])\n",
    "    \n",
    "    # Return the number of elements of the dataset.\n",
    "    def __len__(self):\n",
    "        return len(self.annotations)\n",
    "    \n",
    "# Let's test the data class.\n",
    "trainData = CocoCaptions('annotations/captions_train2014.json')\n",
    "print('Number of training examples: ', len(trainData))\n",
    "\n",
    "# It would be a mistake to build a vocabulary using the validation set so we reuse.\n",
    "valData = CocoCaptions('annotations/captions_val2014.json', vocabulary = trainData.vocabulary)\n",
    "print('Number of validation examples: ', len(valData))\n",
    "\n",
    "# Print a sample from the training data.\n",
    "imgId, caption = trainData[0]\n",
    "print('imgId', imgId)\n",
    "print('caption', caption.tolist())\n",
    "print('captionString', trainData.ids2caption(caption))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Making a Data Loader that can Handle Sequences.\n",
    "\n",
    "Handling sequences is special when processing batches of inputs because each sequence can have a different length. This makes batching complicated, and different libraries have different ways of handling this which might be easier or harder to deal with. Here we are padding the sequences to the maximum sequence length in a given batch. Additionally pytorch has nice utility functions that require sorting the sequences in a batch from longest to shortest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The batch builder will pack all sequences of different length into a single tensor by \n",
    "# padding shorter sequences with a padding token.\n",
    "def customBatchBuilder(samples):\n",
    "    imgIds, captionSeqs = zip(*samples)\n",
    "    \n",
    "    # Sort sequences based on length.\n",
    "    seqLengths = [len(seq) for seq in captionSeqs]\n",
    "    maxSeqLength = max(seqLengths)\n",
    "    sorted_list = sorted(zip(list(imgIds), captionSeqs, seqLengths), key = lambda x: -x[2])\n",
    "    imgIds, captionSeqs, seqLengths = zip(*sorted_list)\n",
    "    \n",
    "    # Create tensor with padded sequences.\n",
    "    paddedSeqs = torch.LongTensor(len(imgIds), maxSeqLength)\n",
    "    paddedSeqs.fill_(0)\n",
    "    for (i, seq) in enumerate(captionSeqs):\n",
    "        paddedSeqs[i, :len(seq)] = seq\n",
    "    return imgIds, paddedSeqs.t(), seqLengths\n",
    "\n",
    "# Data loaders in pytorch can use a custom batch builder, which we are using here.\n",
    "trainLoader = data.DataLoader(trainData, batch_size = 128, \n",
    "                              shuffle = True, num_workers = 0,\n",
    "                              collate_fn = customBatchBuilder)\n",
    "valLoader = data.DataLoader(valData, batch_size = 128, \n",
    "                            shuffle = False, num_workers = 0,\n",
    "                            collate_fn = customBatchBuilder)\n",
    "\n",
    "# Now let's try using the data loader.\n",
    "index, (imgIds, paddedSeqs, seqLengths) = next(enumerate(trainLoader))\n",
    "print('imgIds', imgIds)\n",
    "print('paddedSequences', paddedSeqs.size())\n",
    "print('seqLengths', seqLengths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## 3. Building our model using a Recurrent Neural Network.\n",
    "We will build a model that predicts the next word based on the previous word using a recurrent neural network. Additionally we will be using an Embedding layer which will assign a unique vector to each word. The network will be trained with a softmax + negative log likelihood loss. Similar to classification we will be trying to optimize for the correct word at each time-step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# By now, we should know that pytorch has a functional implementation (as opposed to class version)\n",
    "# of many common layers, which is especially useful for layers that do not have any parameters.\n",
    "# e.g. relu, sigmoid, softmax, etc.\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class TextGeneratorModel(nn.Module):\n",
    "    # The model has three layers: \n",
    "    #    1. An Embedding layer that turns a sequence of word ids into \n",
    "    #       a sequence of vectors of fixed size: embeddingSize.\n",
    "    #    2. An RNN layer that turns the sequence of embedding vectors into \n",
    "    #       a sequence of hiddenStates.\n",
    "    #    3. A classification layer that turns a sequence of hidden states into a \n",
    "    #       sequence of softmax outputs.\n",
    "    def __init__(self, vocabularySize):\n",
    "        super(TextGeneratorModel, self).__init__()\n",
    "        # See documentation for nn.Embedding here:\n",
    "        # http://pytorch.org/docs/master/nn.html#torch.nn.Embedding\n",
    "        self.embedder = nn.Embedding(vocabularySize, 300)\n",
    "        self.rnn = nn.RNN(300, 512, batch_first = False)\n",
    "        self.classifier = nn.Linear(512, vocabularySize)\n",
    "        self.vocabularySize = vocabularySize\n",
    "\n",
    "    # The forward pass makes the sequences go through the three layers defined above.\n",
    "    def forward(self, paddedSeqs, initialHiddenState):\n",
    "        batchSequenceLength = paddedSeqs.size(0)  # 0-dim is sequence-length-dim.\n",
    "        batchSize = paddedSeqs.size(1)  # 1-dim is batch dimension.\n",
    "        \n",
    "        # Transform word ids into an embedding vector.\n",
    "        embeddingVectors = self.embedder(paddedSeqs)\n",
    "        \n",
    "        # Pass the sequence of word embeddings to the RNN.\n",
    "        rnnOutput, finalHiddenState = self.rnn(embeddingVectors, initialHiddenState)\n",
    "        \n",
    "        # Collapse the batch and sequence-length dimensions in order to use nn.Linear.\n",
    "        flatSeqOutput = rnnOutput.view(-1, 512)\n",
    "        predictions = self.classifier(flatSeqOutput)\n",
    "        \n",
    "        # Expand back the batch and sequence-length dimensions and return. \n",
    "        return predictions.view(batchSequenceLength, batchSize, self.vocabularySize), \\\n",
    "               finalHiddenState\n",
    "\n",
    "# Let's test the model on some input batch.\n",
    "vocabularySize = len(trainData.vocabulary['word2id'])\n",
    "model = TextGeneratorModel(vocabularySize)\n",
    "\n",
    "# Create the initial hidden state for the RNN.\n",
    "index, (imgIds, paddedSeqs, seqLengths) = next(enumerate(trainLoader))\n",
    "initialHiddenState = Variable(torch.Tensor(1, paddedSeqs.size(1), 512).zero_())\n",
    "predictions, _ = model(torch.autograd.Variable(paddedSeqs), initialHiddenState)\n",
    "\n",
    "print('Here are input and output size tensor sizes:')\n",
    "# Inputs are seqLength x batchSize x 1 \n",
    "print('inputs', paddedSeqs.size()) # 10 input sequences.\n",
    "# Outputs are seqLength x batchSize x vocabularySize\n",
    "print('outputs', predictions.size()) # 10 output softmax predictions over our vocabularySize outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Sampling a New Sentence from the Model.\n",
    "\n",
    "The code below uses the RNN network as an RNN cell where we only pass one single input word, and a hidden state vector. Then we keep passing the previously predicted word, and previously predicted hidden state to predict the next word. Since the given model is not trained, it will just output a random sequence of words for now. Ideally, the trained model should also learn when to [END] a sentence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_sentence(model, use_cuda = False):\n",
    "    counter = 0\n",
    "    limit = 200\n",
    "    words = list()\n",
    "\n",
    "    # Setup initial input state, and input word (we use \"the\").\n",
    "    previousWord = torch.LongTensor(1, 1).fill_(trainData.vocabulary['word2id']['the'])\n",
    "    previousHiddenState = torch.autograd.Variable(torch.Tensor(1, 1, 512).zero_())\n",
    "    if use_cuda: previousHiddenState = previousHiddenState.cuda()\n",
    "\n",
    "    while True:\n",
    "        # Predict the next word based on the previous hidden state and previous word.\n",
    "        inputWord = torch.autograd.Variable(previousWord)\n",
    "        if use_cuda: inputWord = inputWord.cuda()\n",
    "        predictions, hiddenState = model(inputWord, previousHiddenState)\n",
    "        nextWordId = np.random.multinomial(1, F.softmax(predictions.squeeze()).data.cpu().numpy(), 1).argmax()\n",
    "        words.append(trainData.vocabulary['id2word'][nextWordId])\n",
    "        # Setup the inputs for the next round.\n",
    "        previousWord.fill_(nextWordId)\n",
    "        previousHiddenState = hiddenState\n",
    "\n",
    "        # Keep adding words until the [END] token is generated.\n",
    "        if nextWordId == trainData.vocabulary['word2id']['[END]'] or counter > limit:\n",
    "            break\n",
    "        counter += 1\n",
    "    \n",
    "    words.insert(0, 'the')\n",
    "    words.insert(0, '[START]')\n",
    "    return string.join(words, \" \")\n",
    "\n",
    "print(sample_sentence(model, use_cuda = False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Training the Model\n",
    "\n",
    "Now that data is pre-processed, we can try training the model. An important part is to define our target labels or ground-truth labels. In this text generation model, we want to predict the next word based on the previous word. So we need to provide as the target a shifted version of the input sequence. The code below looks a lot like the code used for training previous models with only small modifications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tqdm as tqdmx\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "tqdmx.tqdm.get_lock().locks = []\n",
    "\n",
    "def train_rnn_model(model, criterion, optimizer, trainLoader, valLoader, n_epochs = 10, use_gpu = False):\n",
    "    return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to the actual training call, notice how unlike previous experiments we are using here RMSprop which is a different type of optimizer that is often preferred for recurrent neural networks, although others such as SGD, and ADAM will also work. Additionally we are using nn.NLLLoss for the loss function, which is equivalent to the nn.CrossEntropyLoss function used before. The only difference is that nn.CrossEntropyLoss does the log_softmax operation for us, however in our implementation, we already applied log_softmax to the outputs of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocabularySize = len(trainData.vocabulary['word2id'])\n",
    "model = TextGeneratorModel(vocabularySize)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr = 0.0005)\n",
    "\n",
    "# Train the previously defined model.\n",
    "train_rnn_model(model, criterion, optimizer, trainLoader, valLoader, n_epochs = 10, use_gpu = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lab Questions (10pts)\n",
    "<span><b>1.</b></span> (2pts) What is the number of parameters of the TextGeneratorModel? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Show how did you come up with that number here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span><b>2.</b></span> (4pts) Provide an implementation for the function train_rnn_model from section 3, this will be similar to the train_model function used in the previous lab. Then train the model and report a few sentences generated by your model. Use the following figure as reference to make sure you are using the right inputs and targets to train the model. The loss function between predictions and targets should be nn.CrossEntropyLoss(), so you might need to collapse the batch and sequence-length dimensions before passing them to the loss function.\n",
    "\n",
    "<img src=\"rnn.png\" width=\"80%\"> "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# implement train_rnn_model and then train the model using this function. \n",
    "# Show here a couple of sentences sampled from your model.\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span><b>3. </b></span> (4pts) Create an ImageCaptioningModel class here that predicts a sentence given an input image. This should be an implementation of the model in this paper https://arxiv.org/pdf/1411.4555.pdf (See figure 3 in the paper). This model is very similar to the one implemented in this lab except that the first RNN cell gets the output of a CNN as its input. I'm also illustrating it below using a figure similar to the one in the previous question. For the CNN use Resnet-18. Note: You do not need to train this model, only define it. Feel free to start from the code for the TextGeneratorModel. <img src=\"im2text.png\" width=\"80%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional Questions (8pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span><b>1. </b></span> (1pts) What is the number of parameters of the ImageCaptioningModel from Q3?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Show how did you come up with that number here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span><b>2. </b></span> (3pts) Modify the TextGeneratorModel to use an LSTM instead, and retrain the model. Report results using this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))\n",
    "print(sample_sentence(model, use_cuda = True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span><b>3. </b></span> (4pts) In this question, you will have to reconstruct an input image from its activations. I will not provide you with the image, only the activation values obtained for a certain layer. You will have access to the code that was used to compute these activations. You will have to use back-propagation to reconstruct the input image. Show the reconstructed input image and tell us who is in the picture. Note: Look at the content reconstruction from outputs performed in https://www.cv-foundation.org/openaccess/content_cvpr_2016/html/Gatys_Image_Style_Transfer_CVPR_2016_paper.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "model = models.vgg16(pretrained = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('image.size()', torch.Size([1, 3, 224, 224]))\n",
      "('layer-19-output.size()', torch.Size([1, 512, 28, 28]))\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image\n",
    "\n",
    "preprocessFn = transforms.Compose([transforms.Scale(256), \n",
    "                                   transforms.CenterCrop(224), \n",
    "                                   transforms.ToTensor(), \n",
    "                                   transforms.Normalize(mean = [0.485, 0.456, 0.406], \n",
    "                                                        std=[0.229, 0.224, 0.225])])\n",
    "\n",
    "def model_F(input, kOutput = 19):\n",
    "    prev_input = input\n",
    "    for layer_id in range(0, kOutput + 1):\n",
    "        current_input = model.features[layer_id](prev_input)\n",
    "        prev_input = current_input\n",
    "    return current_input\n",
    "\n",
    "# Read the incognito image. (Obviously this is not provided in the Lab assignment.)\n",
    "image = preprocessFn(Image.open('incognito.jpg').convert('RGB'))\n",
    "image = Variable(image.unsqueeze(0))\n",
    "\n",
    "# Obtain the output of the VGG layer 19.\n",
    "model.eval()\n",
    "target = Variable(model_F(image).data) # Repack variable.\n",
    "\n",
    "print('image.size()', image.size())\n",
    "print('layer-19-output.size()', target.size())\n",
    "torch.save(target.data, open('layer-19-output.p', 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 512, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "def imshow(img):\n",
    "    # convert torch tensor to PIL image and then show image inline.\n",
    "    img = transforms.ToPILImage()(img[0].cpu() * 0.5 + 0.5) # denormalize tensor before convert\n",
    "    plt.imshow(img, aspect = None)\n",
    "    plt.axis('off')\n",
    "    plt.gcf().set_size_inches(4, 4)\n",
    "    plt.show()\n",
    "\n",
    "target = torch.load(open('layer-19-output.p'))\n",
    "print(target.size())\n",
    "\n",
    "# Your solution goes here. Show the reconstructed input and tell us who is depicted in the incognito.jpg image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"font-size:0.8em;color:#888;text-align:center;padding-top:20px;\">If you find any errors or omissions in this material please contact me at vicente@virginia.edu</div>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
