Training the CNN + RNN Model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Training Setup\n",
    "\n",
    "### Task #1\n",
    "\n",
    "Setting the following variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n",
      "Vocabulary successfully loaded from vocab.pkl file!\n",
      "loading annotations into memory...\n",
      "Done (t=1.04s)\n",
      "creating index...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 759/414113 [00:00<01:54, 3618.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "index created!\n",
      "Obtaining caption lengths...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 414113/414113 [01:34<00:00, 4373.05it/s]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms\n",
    "import sys\n",
    "sys.path.append('/opt/cocoapi/PythonAPI')\n",
    "from pycocotools.coco import COCO\n",
    "from data_loader import get_loader\n",
    "from model import EncoderCNN, DecoderRNN\n",
    "import math\n",
    "import nltk\n",
    "nltk.download('punkt')\n",
    "\n",
    "batch_size = 128         # batch size\n",
    "vocab_threshold = 7        # minimum word count threshold\n",
    "vocab_from_file = True   # if True, load existing vocab file\n",
    "embed_size = 256           # dimensionality of image and word embeddings\n",
    "hidden_size = 512          # number of features in hidden state of the RNN decoder\n",
    "num_epochs = 3             # number of training epochs\n",
    "save_every = 1             # determines frequency of saving model weights\n",
    "print_every = 100          # determines window for printing average loss\n",
    "log_file = 'training_log.txt'       # name of file with saved training loss and perplexity\n",
    "\n",
    "# writing image transformations\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",
    "# Build 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=vocab_from_file)\n",
    "\n",
    "# The size of the vocabulary.\n",
    "vocab_size = len(data_loader.dataset.vocab)\n",
    "\n",
    "# Initialize the encoder and decoder. \n",
    "encoder = EncoderCNN(embed_size)\n",
    "decoder = DecoderRNN(embed_size, hidden_size, vocab_size)\n",
    "\n",
    "# Move models to GPU if CUDA is available. \n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "encoder.to(device)\n",
    "decoder.to(device)\n",
    "\n",
    "# Define the loss function. \n",
    "criterion = nn.CrossEntropyLoss().cuda() if torch.cuda.is_available() else nn.CrossEntropyLoss()\n",
    "\n",
    "# Specify the learnable parameters of the model.\n",
    "params = list(decoder.parameters()) + list(encoder.embed.parameters())\n",
    "          \n",
    "# Define the optimizer.\n",
    "optimizer = torch.optim.Adam(params, lr=0.001)\n",
    "\n",
    "# Set the total number of training steps per epoch.\n",
    "total_step = math.ceil(len(data_loader.dataset.caption_lengths) / data_loader.batch_sampler.batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Train your Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/3], Step [100/3236], Loss: 2.0393, Perplexity: 7.6855\n",
      "Epoch [1/3], Step [200/3236], Loss: 2.0037, Perplexity: 7.41624\n",
      "Epoch [1/3], Step [300/3236], Loss: 2.0837, Perplexity: 8.03388\n",
      "Epoch [1/3], Step [400/3236], Loss: 2.2009, Perplexity: 9.03354\n",
      "Epoch [1/3], Step [500/3236], Loss: 2.1110, Perplexity: 8.25680\n",
      "Epoch [1/3], Step [600/3236], Loss: 2.1678, Perplexity: 8.73933\n",
      "Epoch [1/3], Step [700/3236], Loss: 2.0030, Perplexity: 7.41118\n",
      "Epoch [1/3], Step [800/3236], Loss: 2.1789, Perplexity: 8.83651\n",
      "Epoch [1/3], Step [900/3236], Loss: 2.3058, Perplexity: 10.0317\n",
      "Epoch [1/3], Step [1000/3236], Loss: 1.9142, Perplexity: 6.7816\n",
      "Epoch [1/3], Step [1100/3236], Loss: 1.9716, Perplexity: 7.18253\n",
      "Epoch [1/3], Step [1200/3236], Loss: 2.0434, Perplexity: 7.71652\n",
      "Epoch [1/3], Step [1300/3236], Loss: 1.9894, Perplexity: 7.31139\n",
      "Epoch [1/3], Step [1400/3236], Loss: 2.3539, Perplexity: 10.5265\n",
      "Epoch [1/3], Step [1500/3236], Loss: 1.9957, Perplexity: 7.35769\n",
      "Epoch [1/3], Step [1600/3236], Loss: 2.2887, Perplexity: 9.86222\n",
      "Epoch [1/3], Step [1700/3236], Loss: 2.0040, Perplexity: 7.41893\n",
      "Epoch [1/3], Step [1800/3236], Loss: 1.9025, Perplexity: 6.70271\n",
      "Epoch [1/3], Step [1900/3236], Loss: 1.8822, Perplexity: 6.56810\n",
      "Epoch [1/3], Step [2000/3236], Loss: 2.0622, Perplexity: 7.86307\n",
      "Epoch [1/3], Step [2100/3236], Loss: 1.8720, Perplexity: 6.50140\n",
      "Epoch [1/3], Step [2200/3236], Loss: 2.2023, Perplexity: 9.04583\n",
      "Epoch [1/3], Step [2300/3236], Loss: 2.1674, Perplexity: 8.73529\n",
      "Epoch [1/3], Step [2400/3236], Loss: 1.9473, Perplexity: 7.00948\n",
      "Epoch [1/3], Step [2500/3236], Loss: 2.7885, Perplexity: 16.2569\n",
      "Epoch [1/3], Step [2600/3236], Loss: 1.8777, Perplexity: 6.53832\n",
      "Epoch [1/3], Step [2700/3236], Loss: 2.0274, Perplexity: 7.59430\n",
      "Epoch [1/3], Step [2800/3236], Loss: 1.9749, Perplexity: 7.20615\n",
      "Epoch [1/3], Step [2900/3236], Loss: 1.9760, Perplexity: 7.21419\n",
      "Epoch [1/3], Step [3000/3236], Loss: 2.0106, Perplexity: 7.46804\n",
      "Epoch [1/3], Step [3100/3236], Loss: 1.9717, Perplexity: 7.18316\n",
      "Epoch [1/3], Step [3200/3236], Loss: 2.1337, Perplexity: 8.44641\n",
      "Epoch [2/3], Step [100/3236], Loss: 1.8344, Perplexity: 6.261401\n",
      "Epoch [2/3], Step [200/3236], Loss: 1.9551, Perplexity: 7.06449\n",
      "Epoch [2/3], Step [300/3236], Loss: 1.9761, Perplexity: 7.21487\n",
      "Epoch [2/3], Step [400/3236], Loss: 1.9204, Perplexity: 6.82369\n",
      "Epoch [2/3], Step [500/3236], Loss: 2.5438, Perplexity: 12.7282\n",
      "Epoch [2/3], Step [600/3236], Loss: 2.2102, Perplexity: 9.11753\n",
      "Epoch [2/3], Step [700/3236], Loss: 1.7914, Perplexity: 5.99816\n",
      "Epoch [2/3], Step [800/3236], Loss: 1.8178, Perplexity: 6.15812\n",
      "Epoch [2/3], Step [900/3236], Loss: 2.0820, Perplexity: 8.02039\n",
      "Epoch [2/3], Step [1000/3236], Loss: 2.0786, Perplexity: 7.9930\n",
      "Epoch [2/3], Step [1100/3236], Loss: 1.8634, Perplexity: 6.44552\n",
      "Epoch [2/3], Step [1200/3236], Loss: 1.9886, Perplexity: 7.30529\n",
      "Epoch [2/3], Step [1300/3236], Loss: 1.8705, Perplexity: 6.49162\n",
      "Epoch [2/3], Step [1400/3236], Loss: 1.8621, Perplexity: 6.43750\n",
      "Epoch [2/3], Step [1500/3236], Loss: 1.9245, Perplexity: 6.85179\n",
      "Epoch [2/3], Step [1600/3236], Loss: 1.8070, Perplexity: 6.09217\n",
      "Epoch [2/3], Step [1700/3236], Loss: 1.8529, Perplexity: 6.37806\n",
      "Epoch [2/3], Step [1800/3236], Loss: 2.3940, Perplexity: 10.9569\n",
      "Epoch [2/3], Step [1900/3236], Loss: 1.9118, Perplexity: 6.76503\n",
      "Epoch [2/3], Step [2000/3236], Loss: 1.9764, Perplexity: 7.21643\n",
      "Epoch [2/3], Step [2100/3236], Loss: 2.5962, Perplexity: 13.4131\n",
      "Epoch [2/3], Step [2200/3236], Loss: 2.4606, Perplexity: 11.7120\n",
      "Epoch [2/3], Step [2300/3236], Loss: 1.9472, Perplexity: 7.00872\n",
      "Epoch [2/3], Step [2400/3236], Loss: 1.9294, Perplexity: 6.88561\n",
      "Epoch [2/3], Step [2500/3236], Loss: 3.5868, Perplexity: 36.1201\n",
      "Epoch [2/3], Step [2600/3236], Loss: 2.1147, Perplexity: 8.28713\n",
      "Epoch [2/3], Step [2700/3236], Loss: 1.7557, Perplexity: 5.78756\n",
      "Epoch [2/3], Step [2800/3236], Loss: 1.8051, Perplexity: 6.08063\n",
      "Epoch [2/3], Step [2900/3236], Loss: 1.8938, Perplexity: 6.64480\n",
      "Epoch [2/3], Step [3000/3236], Loss: 1.8904, Perplexity: 6.62222\n",
      "Epoch [2/3], Step [3100/3236], Loss: 1.9451, Perplexity: 6.99419\n",
      "Epoch [2/3], Step [3200/3236], Loss: 1.8358, Perplexity: 6.27057\n",
      "Epoch [3/3], Step [100/3236], Loss: 1.8486, Perplexity: 6.350779\n",
      "Epoch [3/3], Step [200/3236], Loss: 1.8276, Perplexity: 6.21907\n",
      "Epoch [3/3], Step [300/3236], Loss: 1.8740, Perplexity: 6.51453\n",
      "Epoch [3/3], Step [400/3236], Loss: 1.8461, Perplexity: 6.33515\n",
      "Epoch [3/3], Step [500/3236], Loss: 1.8411, Perplexity: 6.30379\n",
      "Epoch [3/3], Step [600/3236], Loss: 1.8993, Perplexity: 6.68116\n",
      "Epoch [3/3], Step [700/3236], Loss: 1.7979, Perplexity: 6.03711\n",
      "Epoch [3/3], Step [800/3236], Loss: 1.8720, Perplexity: 6.50100\n",
      "Epoch [3/3], Step [900/3236], Loss: 1.7748, Perplexity: 5.89893\n",
      "Epoch [3/3], Step [1000/3236], Loss: 1.8434, Perplexity: 6.3177\n",
      "Epoch [3/3], Step [1100/3236], Loss: 1.8287, Perplexity: 6.22581\n",
      "Epoch [3/3], Step [1200/3236], Loss: 1.8255, Perplexity: 6.20594\n",
      "Epoch [3/3], Step [1300/3236], Loss: 1.9524, Perplexity: 7.04567\n",
      "Epoch [3/3], Step [1400/3236], Loss: 1.7988, Perplexity: 6.04225\n",
      "Epoch [3/3], Step [1500/3236], Loss: 1.9713, Perplexity: 7.18014\n",
      "Epoch [3/3], Step [1600/3236], Loss: 2.0450, Perplexity: 7.72958\n",
      "Epoch [3/3], Step [1700/3236], Loss: 1.8698, Perplexity: 6.48678\n",
      "Epoch [3/3], Step [1800/3236], Loss: 1.7946, Perplexity: 6.01731\n",
      "Epoch [3/3], Step [1900/3236], Loss: 1.7594, Perplexity: 5.80900\n",
      "Epoch [3/3], Step [2000/3236], Loss: 1.9765, Perplexity: 7.21761\n",
      "Epoch [3/3], Step [2100/3236], Loss: 1.8626, Perplexity: 6.44026\n",
      "Epoch [3/3], Step [2200/3236], Loss: 2.2390, Perplexity: 9.38419\n",
      "Epoch [3/3], Step [2300/3236], Loss: 1.8710, Perplexity: 6.49481\n",
      "Epoch [3/3], Step [2400/3236], Loss: 1.8341, Perplexity: 6.25950\n",
      "Epoch [3/3], Step [2500/3236], Loss: 1.8793, Perplexity: 6.54864\n",
      "Epoch [3/3], Step [2600/3236], Loss: 1.7952, Perplexity: 6.02100\n",
      "Epoch [3/3], Step [2700/3236], Loss: 1.6823, Perplexity: 5.37800\n",
      "Epoch [3/3], Step [2800/3236], Loss: 1.9075, Perplexity: 6.73600\n",
      "Epoch [3/3], Step [2900/3236], Loss: 1.8130, Perplexity: 6.12919\n",
      "Epoch [3/3], Step [3000/3236], Loss: 1.7719, Perplexity: 5.88202\n",
      "Epoch [3/3], Step [3100/3236], Loss: 1.7932, Perplexity: 6.00883\n",
      "Epoch [3/3], Step [3200/3236], Loss: 1.8262, Perplexity: 6.21033\n",
      "Epoch [3/3], Step [3236/3236], Loss: 1.9442, Perplexity: 6.98773"
     ]
    }
   ],
   "source": [
    "from workspace_utils import active_session\n",
    "import torch.utils.data as data\n",
    "import numpy as np\n",
    "import os\n",
    "import requests\n",
    "import time\n",
    "\n",
    "# Open the training log file.\n",
    "f = open(log_file, 'w')\n",
    "\n",
    "with active_session():\n",
    "    for epoch in range(1, num_epochs+1):\n",
    "\n",
    "        for i_step in range(1, total_step+1):\n",
    "\n",
    "            # Randomly sample a caption length, and sample indices with that length.\n",
    "            indices = data_loader.dataset.get_train_indices()\n",
    "            # Create and assign a batch sampler to retrieve a batch with the sampled indices.\n",
    "            new_sampler = data.sampler.SubsetRandomSampler(indices=indices)\n",
    "            data_loader.batch_sampler.sampler = new_sampler\n",
    "\n",
    "            # Obtain the batch.\n",
    "            images, captions = next(iter(data_loader))\n",
    "\n",
    "            # Move batch of images and captions to GPU if CUDA is available.\n",
    "            images = images.to(device)\n",
    "            captions = captions.to(device)\n",
    "\n",
    "            # Zero the gradients.\n",
    "            decoder.zero_grad()\n",
    "            encoder.zero_grad()\n",
    "\n",
    "            # Pass the inputs through the CNN-RNN model.\n",
    "            features = encoder(images)\n",
    "            outputs = decoder(features, captions)\n",
    "\n",
    "            # Calculate the batch loss.\n",
    "            loss = criterion(outputs.view(-1, vocab_size), captions.view(-1))\n",
    "\n",
    "            # Backward pass.\n",
    "            loss.backward()\n",
    "\n",
    "            # Update the parameters in the optimizer.\n",
    "            optimizer.step()\n",
    "\n",
    "            # Get training statistics.\n",
    "            stats = 'Epoch [%d/%d], Step [%d/%d], Loss: %.4f, Perplexity: %5.4f' % (epoch, num_epochs, i_step, total_step, loss.item(), np.exp(loss.item()))\n",
    "\n",
    "            # Print training statistics (on same line).\n",
    "            print('\\r' + stats, end=\"\")\n",
    "            sys.stdout.flush()\n",
    "\n",
    "            # Print training statistics to file.\n",
    "            f.write(stats + '\\n')\n",
    "            f.flush()\n",
    "\n",
    "            # Print training statistics (on different line).\n",
    "            if i_step % print_every == 0:\n",
    "                print('\\r' + stats)\n",
    "\n",
    "        # Save the weights.\n",
    "        if epoch % save_every == 0:\n",
    "            torch.save(decoder.state_dict(), os.path.join('./models', 'decoder-%d.pkl' % epoch))\n",
    "            torch.save(encoder.state_dict(), os.path.join('./models', 'encoder-%d.pkl' % epoch))\n",
    "\n",
    "# Close the training log file.\n",
    "f.close()"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}