pix2code-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# \"pix2code\" (PyTorch Implementation)\n",
    "#### Version: *alpha (wip)*\n",
    "\n",
    "This is a PyTorch-based implementation of the [work done by Tony Beltramelli on pix2code](https://arxiv.org/abs/1705.07962), using the Bootstrap/DSL dataset created for the paper. The implementation is an image-captioning pair of encoder and decoder models that use the feature extraction of ResNet-152 as a base.\n",
    "\n",
    "Heavily influenced by these to get to a working prototype:\n",
    "- [PyTorch tutorial on image captioning (GitHub)](https://github.com/yunjey/pytorch-tutorial/tree/master/tutorials/03-advanced/image_captioning)\n",
    "- [FloydHub blog post on using keras to transform screenshots to code](https://blog.floydhub.com/Turning-design-mockups-into-code-with-deep-learning/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pdb\n",
    "import os\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.autograd import Variable\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import datasets, models, transforms\n",
    "import numpy as np\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparams\n",
    "batch_size = 4\n",
    "embed_size = 256\n",
    "num_epochs = 1000\n",
    "learning_rate = 0.001\n",
    "hidden_size = 512\n",
    "num_layers = 1\n",
    "\n",
    "# Other params\n",
    "shuffle = True\n",
    "num_workers = 2\n",
    "\n",
    "# Logging Variables\n",
    "save_after_x_epochs = 50\n",
    "log_step = 5\n",
    "\n",
    "# Paths\n",
    "data_dir = './processed_data/data_train/' # For testing purposes, we use a pre-split dataset rather than do it here\n",
    "model_path = './models/'\n",
    "vocab_path = './bootstrap.vocab'\n",
    "\n",
    "# DO NOT CHANGE:\n",
    "crop_size = 224 # Required by resnet152"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building Vocabulary\n",
    "We use the provided DSL vocabulary from pix2code's dataset, consisting of 18 tokens, each of which map to Bootstrap-based HTML code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_doc(filename):\n",
    "    file = open(filename, 'r')\n",
    "    text = file.read()\n",
    "    file.close()\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocabulary (object):\n",
    "    def __init__ (self):\n",
    "        self.word2idx = {}\n",
    "        self.idx2word = {}\n",
    "        self.idx = 0\n",
    "        \n",
    "    def add_word (self, word):\n",
    "        if not word in self.word2idx:\n",
    "            self.word2idx[word] = self.idx\n",
    "            self.idx2word[self.idx] = word\n",
    "            self.idx += 1\n",
    "    \n",
    "    def __call__ (self, word):\n",
    "        if not word in self.word2idx:\n",
    "            return self.word2idx['<unk>']\n",
    "        return self.word2idx[word]\n",
    "    \n",
    "    def __len__ (self):\n",
    "        return len(self.word2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_vocab (vocab_file_path):\n",
    "    vocab = Vocabulary()\n",
    "\n",
    "    # Load the vocab file (super basic split())\n",
    "    words_raw = load_doc(vocab_file_path)\n",
    "    words = set(words_raw.split(' '))\n",
    "    \n",
    "    for i, word in enumerate(words):\n",
    "        vocab.add_word(word)\n",
    "\n",
    "    vocab.add_word(' ')\n",
    "    vocab.add_word('<unk>') # If we find an unknown word\n",
    "    \n",
    "    print('Created vocabulary of ' + str(len(vocab)) + ' items from ' + vocab_file_path)\n",
    "    \n",
    "    return vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created vocabulary of 19 items from ./bootstrap.vocab\n"
     ]
    }
   ],
   "source": [
    "# Load vocabulary\n",
    "vocab = build_vocab(vocab_path)\n",
    "\n",
    "vocab_size = len(vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Dataset (images and captions)\n",
    "Due to the way the dataset was provided ([.gui and .png files in the same folder](https://github.com/tonybeltramelli/pix2code/tree/master/datasets)), we create a custom PyTorch dataloader which stores captions in memory, but loads images on-demand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageHTMLDataSet (Dataset):\n",
    "    def __init__ (self, data_dir, vocab, transform):\n",
    "        self.data_dir = data_dir\n",
    "        self.vocab = vocab\n",
    "        self.transform = transform\n",
    "        \n",
    "        self.raw_image_names = []\n",
    "        self.raw_captions = []\n",
    "        \n",
    "        # Fetch all files from our data directoruy\n",
    "        self.filenames = os.listdir(data_dir)\n",
    "        self.filenames.sort()\n",
    "        \n",
    "        # Sort files based on their filetype\n",
    "        # Assume associated training examples have same filenames\n",
    "        for filename in self.filenames:\n",
    "            if filename[-3:] == 'png':\n",
    "                # Store image filename\n",
    "                self.raw_image_names.append(filename)\n",
    "            elif filename[-3:] == 'gui':\n",
    "                # Load .gui file\n",
    "                data = load_doc(data_dir + filename)\n",
    "                self.raw_captions.append(data)\n",
    "                \n",
    "        print('Created dataset of ' + str(len(self)) + ' items from ' + data_dir)\n",
    "\n",
    "    def __len__ (self):\n",
    "        return len(self.raw_image_names)\n",
    "    \n",
    "    def __getitem__ (self, idx):\n",
    "        img_path, raw_caption = self.raw_image_names[idx], self.raw_captions[idx]\n",
    "        \n",
    "        # Get image from filesystem\n",
    "        image = Image.open(os.path.join(self.data_dir, img_path)).convert('RGB')\n",
    "        image = self.transform(image)\n",
    "        \n",
    "        # Convert caption (string) to list of vocab ID's\n",
    "        caption = []\n",
    "        caption.append(self.vocab('<START>'))\n",
    "        \n",
    "        # Remove newlines, separate words with spaces\n",
    "        tokens = ' '.join(raw_caption.split())\n",
    "\n",
    "        # Add space after each comma\n",
    "        tokens = tokens.replace(',', ' ,')\n",
    "        \n",
    "        # Split into words\n",
    "        tokens = tokens.split(' ')\n",
    "        \n",
    "        caption.extend([self.vocab(token) for token in tokens])\n",
    "        caption.append(self.vocab('<END>'))\n",
    "        \n",
    "        target = torch.Tensor(caption)\n",
    "        \n",
    "        return image, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See https://github.com/yunjey/pytorch-tutorial/tree/master/tutorials/03-advanced/image_captioning\n",
    "def collate_fn (data):\n",
    "    # Sort datalist by caption length; descending order\n",
    "    data.sort(key = lambda data_pair: len(data_pair[1]), reverse=True)\n",
    "    images, captions = zip(*data)\n",
    "    \n",
    "    # Merge images (from tuple of 3D Tensor to 4D Tensor)\n",
    "    images = torch.stack(images, 0)\n",
    "    \n",
    "    # Merge captions (from tuple of 1D tensor to 2D tensor)\n",
    "    lengths = [len(caption) for caption in captions] # List of caption lengths\n",
    "    targets = torch.zeros(len(captions), max(lengths)).long()\n",
    "    \n",
    "    for i, caption in enumerate(captions):\n",
    "        end = lengths[i]\n",
    "        targets[i, :end] = caption[:end]\n",
    "        \n",
    "    return images, targets, lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created dataset of 1400 items from ./processed_data/data_train/\n"
     ]
    }
   ],
   "source": [
    "# Transform to modify images for pre-trained ResNet base\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((crop_size, crop_size)), # Match resnet size\n",
    "    transforms.ToTensor(),\n",
    "    # See for magic #'s: http://pytorch.org/docs/master/torchvision/models.html\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# Create data loader\n",
    "img_html_dataset = ImageHTMLDataSet(data_dir=data_dir, vocab=vocab, transform=transform)\n",
    "data_loader = DataLoader(dataset=img_html_dataset,\n",
    "                         batch_size=batch_size,\n",
    "                         shuffle=shuffle,\n",
    "                         num_workers=num_workers,\n",
    "                         collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encoder Model\n",
    "\n",
    "Takes in input matrix of images and outputs features based on ResNet-152."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderCNN (nn.Module):\n",
    "    def __init__ (self, embed_size):\n",
    "        super(EncoderCNN, self).__init__()\n",
    "        \n",
    "        # Load pretrained resnet model\n",
    "        resnet = models.resnet152(pretrained = True)\n",
    "        \n",
    "        # Remove the fully connected layers\n",
    "        modules = list(resnet.children())[:-1]\n",
    "        self.resnet = nn.Sequential(*modules)\n",
    "        \n",
    "        # Create our replacement layers\n",
    "        # We reuse the in_feature size of the resnet fc layer for our first replacement layer = 2048 as of creation\n",
    "        self.linear = nn.Linear(in_features = resnet.fc.in_features, out_features = embed_size)\n",
    "        self.bn = nn.BatchNorm1d(num_features = embed_size, momentum = 0.01)\n",
    "        \n",
    "        print('EncoderCNN created with embed_size: ' + str(embed_size))\n",
    "\n",
    "    def forward (self, images):\n",
    "        # Get the expected output from the fully connected layers\n",
    "        # Fn: AvgPool2d(kernel_size=7, stride=1, padding=0, ceil_mode=False, count_include_pad=True)\n",
    "        # Output: torch.Size([batch_size, 2048, 1, 1])\n",
    "        features = self.resnet(images)\n",
    "\n",
    "        # Resize the features for our linear function\n",
    "        features = features.view(features.size(0), -1)\n",
    "        \n",
    "        # Fn: Linear(in_features=2048, out_features=embed_size, bias=True)\n",
    "        # Output: torch.Size([batch_size, embed_size])\n",
    "        features = self.linear(features)\n",
    "        \n",
    "        # Fn: BatchNorm1d(embed_size, eps=1e-05, momentum=0.01, affine=True)\n",
    "        # Output: torch.Size([batch_size, embed_size])\n",
    "        features = self.bn(features)\n",
    "        \n",
    "        return features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decoder Model\n",
    "We can substitute the LSTM for a GRU to speed up our training (as suggested in the FloydHub post)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderRNN (nn.Module):\n",
    "    def __init__ (self, embed_size, hidden_size, vocab_size, num_layers):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        \n",
    "        # 19 word vocabulary, embed_size dimensional embeddings\n",
    "        self.embed = nn.Embedding(num_embeddings = vocab_size, embedding_dim = embed_size)\n",
    "\n",
    "        self.lstm = nn.LSTM(input_size = embed_size, hidden_size = hidden_size, num_layers = num_layers, batch_first=True)\n",
    "\n",
    "        self.linear = nn.Linear(in_features = hidden_size, out_features = vocab_size)\n",
    "        \n",
    "        # Store the embed size for use when sampling\n",
    "        self.embed_size = embed_size\n",
    "        \n",
    "        print('DecoderRNN created with embed_size: ' + str(embed_size))\n",
    "        \n",
    "    def forward (self, features, captions, lengths):\n",
    "        # 'captions' enters as shape torch.Size([batch_size, len(longest caption)])\n",
    "        \n",
    "        # Fn: Embedding(vocab_size, embed_size)\n",
    "        # Input: LongTensor (N = mini_batch, W = # of indices to extract per mini-batch)\n",
    "        # Output: (N, W, embedding_dim) => torch.Size([batch_size, len(longest caption), embed_size])\n",
    "        embeddings = self.embed(captions)\n",
    "        \n",
    "        # Match features dimensions to embedding's\n",
    "        features = features.unsqueeze(1) # torch.Size([4, 128]) => torch.Size([4, 1, 128])\n",
    "        \n",
    "        embeddings = torch.cat((features, embeddings), 1)\n",
    "        \n",
    "        packed = nn.utils.rnn.pack_padded_sequence(input = embeddings, lengths = lengths, batch_first = True)\n",
    "        \n",
    "        # Fn: LSTM(embed_size, hidden_size, batch_first = True)\n",
    "        hiddens, _ = self.lstm(packed)\n",
    "        \n",
    "        outputs = self.linear(hiddens[0])\n",
    "        \n",
    "        return outputs\n",
    "    \n",
    "    # Sample method used for testing our model\n",
    "    def sample (self, features, states=None):\n",
    "        sampled_ids = []\n",
    "        inputs = features.unsqueeze(1)\n",
    "        \n",
    "        # Put the features input through our decoder for i iterations\n",
    "        # TODO: Put this range into a parameter?\n",
    "        for i in range(100):\n",
    "            hiddens, states = self.lstm(inputs, states)\n",
    "            outputs = self.linear(hiddens.squeeze(1))\n",
    "            predicted = outputs.max(dim = 1, keepdim = True)[1]\n",
    "            sampled_ids.append(predicted)\n",
    "            inputs = self.embed(predicted)\n",
    "            inputs = inputs.view(-1, 1, self.embed_size)\n",
    "\n",
    "        sampled_ids = torch.cat(sampled_ids, 1)\n",
    "\n",
    "        return sampled_ids.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "CUDA activated.\n"
     ]
    }
   ],
   "source": [
    "encoder = EncoderCNN(embed_size)\n",
    "decoder = DecoderRNN(embed_size, hidden_size, vocab_size, num_layers)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters())\n",
    "optimizer = torch.optim.Adam(params, lr = learning_rate)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    encoder.cuda()\n",
    "    decoder.cuda()\n",
    "    print('CUDA activated.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [#0], Loss: 2.9575, Perplexity: 19.2496\n",
      "Epoch [#5], Loss: 0.1309, Perplexity: 1.1399\n",
      "Epoch [#10], Loss: 0.1157, Perplexity: 1.1226\n",
      "Epoch [#15], Loss: 0.1191, Perplexity: 1.1265\n",
      "Epoch [#20], Loss: 0.1145, Perplexity: 1.1213\n",
      "Epoch [#25], Loss: 0.1521, Perplexity: 1.1642\n",
      "Epoch [#30], Loss: 0.1375, Perplexity: 1.1474\n",
      "Epoch [#35], Loss: 0.1226, Perplexity: 1.1305\n",
      "Epoch [#40], Loss: 0.0897, Perplexity: 1.0938\n",
      "Epoch [#45], Loss: 0.0923, Perplexity: 1.0966\n",
      "!!! saving models at epoch: 49\n",
      "Epoch [#50], Loss: 0.0853, Perplexity: 1.0890\n",
      "Epoch [#55], Loss: 0.1202, Perplexity: 1.1277\n",
      "Epoch [#60], Loss: 0.1151, Perplexity: 1.1220\n",
      "Epoch [#65], Loss: 0.1065, Perplexity: 1.1124\n",
      "Epoch [#70], Loss: 0.0645, Perplexity: 1.0667\n",
      "Epoch [#75], Loss: 0.0788, Perplexity: 1.0820\n",
      "Epoch [#80], Loss: 0.0724, Perplexity: 1.0751\n",
      "Epoch [#85], Loss: 0.0657, Perplexity: 1.0679\n",
      "Epoch [#90], Loss: 0.0977, Perplexity: 1.1026\n",
      "Epoch [#95], Loss: 0.0861, Perplexity: 1.0899\n",
      "!!! saving models at epoch: 99\n",
      "Epoch [#100], Loss: 0.0560, Perplexity: 1.0576\n",
      "Epoch [#105], Loss: 0.0684, Perplexity: 1.0708\n",
      "Epoch [#110], Loss: 0.0845, Perplexity: 1.0882\n",
      "Epoch [#115], Loss: 0.1049, Perplexity: 1.1106\n",
      "Epoch [#120], Loss: 0.0621, Perplexity: 1.0641\n",
      "Epoch [#125], Loss: 0.0602, Perplexity: 1.0621\n",
      "Epoch [#130], Loss: 0.0619, Perplexity: 1.0638\n",
      "Epoch [#135], Loss: 0.0614, Perplexity: 1.0633\n",
      "Epoch [#140], Loss: 0.0673, Perplexity: 1.0696\n",
      "Epoch [#145], Loss: 0.0718, Perplexity: 1.0745\n",
      "!!! saving models at epoch: 149\n",
      "Epoch [#150], Loss: 0.0721, Perplexity: 1.0747\n",
      "Epoch [#155], Loss: 0.0528, Perplexity: 1.0543\n",
      "Epoch [#160], Loss: 0.0495, Perplexity: 1.0507\n",
      "Epoch [#165], Loss: 0.0504, Perplexity: 1.0516\n",
      "Epoch [#170], Loss: 0.0411, Perplexity: 1.0420\n",
      "Epoch [#175], Loss: 0.0324, Perplexity: 1.0329\n",
      "Epoch [#180], Loss: 0.0447, Perplexity: 1.0457\n",
      "Epoch [#185], Loss: 0.0288, Perplexity: 1.0293\n",
      "Epoch [#190], Loss: 0.0287, Perplexity: 1.0291\n",
      "Epoch [#195], Loss: 0.0444, Perplexity: 1.0454\n",
      "!!! saving models at epoch: 199\n",
      "Epoch [#200], Loss: 0.0862, Perplexity: 1.0901\n",
      "Epoch [#205], Loss: 0.0246, Perplexity: 1.0249\n",
      "Epoch [#210], Loss: 0.0453, Perplexity: 1.0464\n",
      "Epoch [#215], Loss: 0.0307, Perplexity: 1.0312\n",
      "Epoch [#220], Loss: 0.0353, Perplexity: 1.0359\n",
      "Epoch [#225], Loss: 0.0384, Perplexity: 1.0391\n",
      "Epoch [#230], Loss: 0.0619, Perplexity: 1.0638\n",
      "Epoch [#235], Loss: 0.0139, Perplexity: 1.0140\n",
      "Epoch [#240], Loss: 0.0242, Perplexity: 1.0245\n",
      "Epoch [#245], Loss: 0.0418, Perplexity: 1.0427\n",
      "!!! saving models at epoch: 249\n",
      "Epoch [#250], Loss: 0.0296, Perplexity: 1.0300\n",
      "Epoch [#255], Loss: 0.0663, Perplexity: 1.0686\n",
      "Epoch [#260], Loss: 0.0368, Perplexity: 1.0374\n",
      "Epoch [#265], Loss: 0.0302, Perplexity: 1.0306\n",
      "Epoch [#270], Loss: 0.0266, Perplexity: 1.0270\n",
      "Epoch [#275], Loss: 0.0186, Perplexity: 1.0187\n",
      "Epoch [#280], Loss: 0.0322, Perplexity: 1.0327\n",
      "Epoch [#285], Loss: 0.0364, Perplexity: 1.0371\n",
      "Epoch [#290], Loss: 0.0290, Perplexity: 1.0295\n",
      "Epoch [#295], Loss: 0.0168, Perplexity: 1.0169\n",
      "!!! saving models at epoch: 299\n",
      "Epoch [#300], Loss: 0.0223, Perplexity: 1.0225\n",
      "Epoch [#305], Loss: 0.0160, Perplexity: 1.0161\n",
      "Epoch [#310], Loss: 0.0281, Perplexity: 1.0285\n",
      "Epoch [#315], Loss: 0.0249, Perplexity: 1.0252\n",
      "Epoch [#320], Loss: 0.0267, Perplexity: 1.0270\n",
      "Epoch [#325], Loss: 0.0148, Perplexity: 1.0149\n",
      "Epoch [#330], Loss: 0.0119, Perplexity: 1.0119\n",
      "Epoch [#335], Loss: 0.0053, Perplexity: 1.0053\n",
      "Epoch [#340], Loss: 0.0129, Perplexity: 1.0130\n",
      "Epoch [#345], Loss: 0.0470, Perplexity: 1.0482\n",
      "!!! saving models at epoch: 349\n",
      "Epoch [#350], Loss: 0.0039, Perplexity: 1.0039\n",
      "Epoch [#355], Loss: 0.0058, Perplexity: 1.0058\n",
      "Epoch [#360], Loss: 0.0081, Perplexity: 1.0082\n",
      "Epoch [#365], Loss: 0.0185, Perplexity: 1.0187\n",
      "Epoch [#370], Loss: 0.0155, Perplexity: 1.0157\n",
      "Epoch [#375], Loss: 0.0157, Perplexity: 1.0158\n",
      "Epoch [#380], Loss: 0.0371, Perplexity: 1.0378\n",
      "Epoch [#385], Loss: 0.0221, Perplexity: 1.0223\n",
      "Epoch [#390], Loss: 0.0054, Perplexity: 1.0054\n",
      "Epoch [#395], Loss: 0.0298, Perplexity: 1.0302\n",
      "!!! saving models at epoch: 399\n",
      "Epoch [#400], Loss: 0.0122, Perplexity: 1.0123\n",
      "Epoch [#405], Loss: 0.0040, Perplexity: 1.0040\n",
      "Epoch [#410], Loss: 0.0256, Perplexity: 1.0260\n",
      "Epoch [#415], Loss: 0.0068, Perplexity: 1.0068\n",
      "Epoch [#420], Loss: 0.0237, Perplexity: 1.0240\n",
      "Epoch [#425], Loss: 0.1460, Perplexity: 1.1571\n",
      "Epoch [#430], Loss: 0.0198, Perplexity: 1.0200\n",
      "Epoch [#435], Loss: 0.0174, Perplexity: 1.0175\n",
      "Epoch [#440], Loss: 0.0152, Perplexity: 1.0154\n",
      "Epoch [#445], Loss: 0.0253, Perplexity: 1.0257\n",
      "!!! saving models at epoch: 449\n",
      "Epoch [#450], Loss: 0.0647, Perplexity: 1.0668\n",
      "Epoch [#455], Loss: 0.0021, Perplexity: 1.0021\n",
      "Epoch [#460], Loss: 0.0042, Perplexity: 1.0042\n",
      "Epoch [#465], Loss: 0.0143, Perplexity: 1.0144\n",
      "Epoch [#470], Loss: 0.0038, Perplexity: 1.0038\n",
      "Epoch [#475], Loss: 0.0131, Perplexity: 1.0132\n",
      "Epoch [#480], Loss: 0.0201, Perplexity: 1.0203\n",
      "Epoch [#485], Loss: 0.0133, Perplexity: 1.0134\n",
      "Epoch [#490], Loss: 0.0112, Perplexity: 1.0113\n",
      "Epoch [#495], Loss: 0.0497, Perplexity: 1.0509\n",
      "!!! saving models at epoch: 499\n",
      "Epoch [#500], Loss: 0.0265, Perplexity: 1.0268\n",
      "Epoch [#505], Loss: 0.0172, Perplexity: 1.0173\n",
      "Epoch [#510], Loss: 0.0320, Perplexity: 1.0326\n",
      "Epoch [#515], Loss: 0.0025, Perplexity: 1.0025\n",
      "Epoch [#520], Loss: 0.0305, Perplexity: 1.0310\n",
      "Epoch [#525], Loss: 0.0201, Perplexity: 1.0204\n",
      "Epoch [#530], Loss: 0.0150, Perplexity: 1.0151\n",
      "Epoch [#535], Loss: 0.0246, Perplexity: 1.0249\n",
      "Epoch [#540], Loss: 0.0071, Perplexity: 1.0072\n",
      "Epoch [#545], Loss: 0.0113, Perplexity: 1.0113\n",
      "!!! saving models at epoch: 549\n",
      "Epoch [#550], Loss: 0.0051, Perplexity: 1.0051\n",
      "Epoch [#555], Loss: 0.0038, Perplexity: 1.0038\n",
      "Epoch [#560], Loss: 0.0422, Perplexity: 1.0431\n",
      "Epoch [#565], Loss: 0.0390, Perplexity: 1.0397\n",
      "Epoch [#570], Loss: 0.0108, Perplexity: 1.0108\n",
      "Epoch [#575], Loss: 0.0253, Perplexity: 1.0257\n",
      "Epoch [#580], Loss: 0.0171, Perplexity: 1.0172\n",
      "Epoch [#585], Loss: 0.0049, Perplexity: 1.0050\n",
      "Epoch [#590], Loss: 0.0119, Perplexity: 1.0119\n",
      "Epoch [#595], Loss: 0.0015, Perplexity: 1.0015\n",
      "!!! saving models at epoch: 599\n",
      "Epoch [#600], Loss: 0.0141, Perplexity: 1.0142\n",
      "Epoch [#605], Loss: 0.0613, Perplexity: 1.0632\n",
      "Epoch [#610], Loss: 0.0018, Perplexity: 1.0018\n",
      "Epoch [#615], Loss: 0.0103, Perplexity: 1.0103\n",
      "Epoch [#620], Loss: 0.0038, Perplexity: 1.0038\n",
      "Epoch [#625], Loss: 0.0047, Perplexity: 1.0047\n",
      "Epoch [#630], Loss: 0.0216, Perplexity: 1.0218\n",
      "Epoch [#635], Loss: 0.0262, Perplexity: 1.0265\n",
      "Epoch [#640], Loss: 0.0294, Perplexity: 1.0298\n",
      "Epoch [#645], Loss: 0.0243, Perplexity: 1.0246\n",
      "!!! saving models at epoch: 649\n",
      "Epoch [#650], Loss: 0.0018, Perplexity: 1.0018\n",
      "Epoch [#655], Loss: 0.0024, Perplexity: 1.0024\n",
      "Epoch [#660], Loss: 0.0479, Perplexity: 1.0491\n",
      "Epoch [#665], Loss: 0.0220, Perplexity: 1.0222\n",
      "Epoch [#670], Loss: 0.0212, Perplexity: 1.0214\n",
      "Epoch [#675], Loss: 0.0021, Perplexity: 1.0021\n",
      "Epoch [#680], Loss: 0.0073, Perplexity: 1.0073\n",
      "Epoch [#685], Loss: 0.0052, Perplexity: 1.0052\n",
      "Epoch [#690], Loss: 0.0284, Perplexity: 1.0288\n",
      "Epoch [#695], Loss: 0.0078, Perplexity: 1.0078\n",
      "!!! saving models at epoch: 699\n",
      "Epoch [#700], Loss: 0.0077, Perplexity: 1.0077\n",
      "Epoch [#705], Loss: 0.0101, Perplexity: 1.0102\n",
      "Epoch [#710], Loss: 0.0023, Perplexity: 1.0023\n",
      "Epoch [#715], Loss: 0.0322, Perplexity: 1.0327\n",
      "Epoch [#720], Loss: 0.0195, Perplexity: 1.0197\n",
      "Epoch [#725], Loss: 0.0138, Perplexity: 1.0139\n",
      "Epoch [#730], Loss: 0.0033, Perplexity: 1.0033\n",
      "Epoch [#735], Loss: 0.0118, Perplexity: 1.0119\n",
      "Epoch [#740], Loss: 0.0026, Perplexity: 1.0026\n",
      "Epoch [#745], Loss: 0.0092, Perplexity: 1.0092\n",
      "!!! saving models at epoch: 749\n",
      "Epoch [#750], Loss: 0.0039, Perplexity: 1.0039\n",
      "Epoch [#755], Loss: 0.0063, Perplexity: 1.0063\n",
      "Epoch [#760], Loss: 0.0214, Perplexity: 1.0217\n",
      "Epoch [#765], Loss: 0.0033, Perplexity: 1.0033\n",
      "Epoch [#770], Loss: 0.0208, Perplexity: 1.0210\n",
      "Epoch [#775], Loss: 0.0177, Perplexity: 1.0178\n",
      "Epoch [#780], Loss: 0.0026, Perplexity: 1.0026\n",
      "Epoch [#785], Loss: 0.0026, Perplexity: 1.0026\n",
      "Epoch [#790], Loss: 0.0034, Perplexity: 1.0034\n",
      "Epoch [#795], Loss: 0.0092, Perplexity: 1.0092\n",
      "!!! saving models at epoch: 799\n",
      "Epoch [#800], Loss: 0.0104, Perplexity: 1.0105\n",
      "Epoch [#805], Loss: 0.0058, Perplexity: 1.0058\n",
      "Epoch [#810], Loss: 0.0159, Perplexity: 1.0160\n",
      "Epoch [#815], Loss: 0.0038, Perplexity: 1.0038\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [#820], Loss: 0.0122, Perplexity: 1.0122\n",
      "Epoch [#825], Loss: 0.0054, Perplexity: 1.0054\n",
      "Epoch [#830], Loss: 0.0020, Perplexity: 1.0020\n",
      "Epoch [#835], Loss: 0.0012, Perplexity: 1.0012\n",
      "Epoch [#840], Loss: 0.0347, Perplexity: 1.0353\n",
      "Epoch [#845], Loss: 0.0070, Perplexity: 1.0070\n",
      "!!! saving models at epoch: 849\n",
      "Epoch [#850], Loss: 0.0022, Perplexity: 1.0022\n",
      "Epoch [#855], Loss: 0.0049, Perplexity: 1.0049\n",
      "Epoch [#860], Loss: 0.0111, Perplexity: 1.0111\n",
      "Epoch [#865], Loss: 0.0263, Perplexity: 1.0266\n",
      "Epoch [#870], Loss: 0.0065, Perplexity: 1.0065\n",
      "Epoch [#875], Loss: 0.0114, Perplexity: 1.0115\n",
      "Epoch [#880], Loss: 0.0162, Perplexity: 1.0163\n",
      "Epoch [#885], Loss: 0.0089, Perplexity: 1.0089\n",
      "Epoch [#890], Loss: 0.0117, Perplexity: 1.0117\n",
      "Epoch [#895], Loss: 0.0360, Perplexity: 1.0367\n",
      "!!! saving models at epoch: 899\n",
      "Epoch [#900], Loss: 0.0186, Perplexity: 1.0188\n",
      "Epoch [#905], Loss: 0.0239, Perplexity: 1.0242\n",
      "Epoch [#910], Loss: 0.0126, Perplexity: 1.0127\n",
      "Epoch [#915], Loss: 0.0072, Perplexity: 1.0072\n",
      "Epoch [#920], Loss: 0.0161, Perplexity: 1.0162\n",
      "Epoch [#925], Loss: 0.0013, Perplexity: 1.0013\n",
      "Epoch [#930], Loss: 0.0061, Perplexity: 1.0061\n",
      "Epoch [#935], Loss: 0.0019, Perplexity: 1.0019\n",
      "Epoch [#940], Loss: 0.0036, Perplexity: 1.0036\n",
      "Epoch [#945], Loss: 0.0208, Perplexity: 1.0210\n",
      "!!! saving models at epoch: 949\n",
      "Epoch [#950], Loss: 0.0093, Perplexity: 1.0093\n",
      "Epoch [#955], Loss: 0.0224, Perplexity: 1.0226\n",
      "Epoch [#960], Loss: 0.0046, Perplexity: 1.0046\n",
      "Epoch [#965], Loss: 0.0131, Perplexity: 1.0131\n",
      "Epoch [#970], Loss: 0.0027, Perplexity: 1.0027\n",
      "Epoch [#975], Loss: 0.0054, Perplexity: 1.0055\n",
      "Epoch [#980], Loss: 0.0598, Perplexity: 1.0616\n",
      "Epoch [#985], Loss: 0.0390, Perplexity: 1.0398\n",
      "Epoch [#990], Loss: 0.0417, Perplexity: 1.0426\n",
      "Epoch [#995], Loss: 0.0015, Perplexity: 1.0015\n",
      "!!! saving models at epoch: 999\n",
      "done!\n"
     ]
    }
   ],
   "source": [
    "encoder.train()\n",
    "decoder.train()\n",
    "\n",
    "batch_count = len(data_loader)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    for i, (images, captions, lengths) in enumerate(data_loader):\n",
    "        # Shape: torch.Size([batch_size, 3, crop_size, crop_size])\n",
    "        images = Variable(images.cuda())\n",
    "\n",
    "        # Shape: torch.Size([batch_size, len(longest caption)])\n",
    "        captions = Variable(captions.cuda())\n",
    "\n",
    "        # lengths is a list of how long captions are in descending order (e.g., [77, 77, 75, 25])\n",
    "\n",
    "        # We remove the paddings from captions that are padded and then pack them into a single sequence\n",
    "        # Our data loader's collate_fn adds extra zeros to the end of sequences that are too short\n",
    "        # Shape: torch.Size([sum(lengths)])\n",
    "        targets = nn.utils.rnn.pack_padded_sequence(input = captions, lengths = lengths, batch_first = True)[0]\n",
    "\n",
    "        # Zero out buffers\n",
    "        encoder.zero_grad()\n",
    "        decoder.zero_grad()\n",
    "\n",
    "        # Forward, Backward, and Optimize\n",
    "        features = encoder(images) # Outputs features of torch.Size([batch_size, embed_size])\n",
    "        outputs = decoder(features, captions, lengths)\n",
    "\n",
    "        # CrossEntropyLoss is expecting:\n",
    "        # Input:  (N, C) where C = number of classes\n",
    "        loss = criterion(outputs, targets)\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        if epoch % log_step == 0 and i == 0:\n",
    "            print('Epoch [#%d], Loss: %.4f, Perplexity: %5.4f'\n",
    "                  % (epoch, loss.data[0], np.exp(loss.data[0])))\n",
    "            \n",
    "        if (epoch + 1) % save_after_x_epochs == 0 and i == 0:\n",
    "            # Save our models\n",
    "            print('!!! saving models at epoch: ' + str(epoch))\n",
    "            torch.save(decoder.state_dict(),os.path.join(model_path, 'decoder-%d-%d.pkl' %(epoch+1, i+1)))\n",
    "            torch.save(encoder.state_dict(), os.path.join(model_path, 'encoder-%d-%d.pkl' %(epoch+1, i+1)))\n",
    "            \n",
    "print('done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(encoder.state_dict(), os.path.join(model_path, 'encoder-1000-1.pkl'))\n",
    "torch.save(decoder.state_dict(), os.path.join(model_path, 'decoder-1000-1.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import corpus_bleu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_idx_to_words (input):\n",
    "    sampled_caption = []\n",
    "    \n",
    "    for idx in input:\n",
    "        word = vocab.idx2word[idx]\n",
    "        sampled_caption.append(word)\n",
    "\n",
    "        if word == '<END>':\n",
    "            break\n",
    "\n",
    "    output = ' '.join(sampled_caption[1:-1])\n",
    "\n",
    "    output = output.replace(' ,', ',')\n",
    "\n",
    "    return output.split(' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_data_dir = './processed_data/data_dev/'\n",
    "\n",
    "# Assume format: \"encoder-# epoch-# iter.pkl\"\n",
    "models_to_test = [\n",
    "    '100-1',\n",
    "    '200-1',\n",
    "    '300-1',\n",
    "    '400-1',\n",
    "    '500-1',\n",
    "    '600-1',\n",
    "    '700-1',\n",
    "    '800-1',\n",
    "    '900-1',\n",
    "    '1000-1',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 100-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 200-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 300-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 400-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 500-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 600-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 700-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 800-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 900-1\n",
      "Created dataset of 175 items from ./processed_data/data_dev/\n",
      "EncoderCNN created with embed_size: 256\n",
      "DecoderRNN created with embed_size: 256\n",
      "done with 175 items for model: 1000-1\n"
     ]
    }
   ],
   "source": [
    "bleu_scores = []\n",
    "\n",
    "for model_idx, model_name in enumerate(models_to_test):\n",
    "    encoder_model_path = os.path.join(model_path, 'encoder-{}.pkl'.format(model_name))\n",
    "    decoder_model_path = os.path.join(model_path, 'decoder-{}.pkl'.format(model_name))\n",
    "    \n",
    "    # Create a data loader for our cross validation testing\n",
    "    dev_img_html_dataset = ImageHTMLDataSet(data_dir=dev_data_dir, vocab=vocab, transform=transform)\n",
    "    \n",
    "    dev_data_loader = DataLoader(dataset=dev_img_html_dataset,\n",
    "                             batch_size=batch_size,\n",
    "                             shuffle=shuffle,\n",
    "                             num_workers=num_workers,\n",
    "                             collate_fn=collate_fn)\n",
    "    \n",
    "    # Load trained models\n",
    "    dev_encoder = EncoderCNN(embed_size)\n",
    "    dev_decoder = DecoderRNN(embed_size, hidden_size, len(vocab), num_layers)\n",
    "\n",
    "    dev_encoder.load_state_dict(torch.load(encoder_model_path))\n",
    "    dev_decoder.load_state_dict(torch.load(decoder_model_path))\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        dev_encoder.cuda()\n",
    "        dev_decoder.cuda()\n",
    "\n",
    "    dev_encoder.eval()\n",
    "    dev_decoder.eval()\n",
    "    \n",
    "    dev_data_count = len(dev_data_loader.dataset)\n",
    "\n",
    "    predicted, actual = list(), list()\n",
    "\n",
    "    for i in range(dev_data_count):\n",
    "        image, caption = dev_data_loader.dataset.__getitem__(i)\n",
    "        image_tensor = Variable(image.unsqueeze(0).cuda())\n",
    "\n",
    "        features = dev_encoder(image_tensor)\n",
    "\n",
    "        sampled_ids = dev_decoder.sample(features)\n",
    "        sampled_ids = sampled_ids.cpu().data.numpy()\n",
    "\n",
    "        predicted.append(sampled_ids)\n",
    "        actual.append(caption.numpy())\n",
    "\n",
    "    predicted = [transform_idx_to_words(item) for item in predicted]\n",
    "    actual = [[transform_idx_to_words(item)] for item in actual]\n",
    "    \n",
    "    bleu = corpus_bleu(actual, predicted)\n",
    "    \n",
    "    bleu_scores.append((model_name, bleu))\n",
    "\n",
    "    print('done with {} items for model: {}'.format(str(len(predicted)), model_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('100-1', 0.924819024813575),\n",
       " ('200-1', 0.8959682936793072),\n",
       " ('300-1', 0.920018035108614),\n",
       " ('400-1', 0.9033686688876422),\n",
       " ('500-1', 0.9191831702256483),\n",
       " ('600-1', 0.8876589013117365),\n",
       " ('700-1', 0.9098200684974314),\n",
       " ('800-1', 0.9164715424716915),\n",
       " ('900-1', 0.919587909269687),\n",
       " ('1000-1', 0.8644991775415475)]"
      ]
     },
     "execution_count": 200,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bleu_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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
}