test_rnn_attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "import torch.backends.cudnn as cudnn\n",
    "import numpy as np\n",
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "from tqdm import tqdm\n",
    "from datasets import *\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "data_folder = '/data3/zhangweiyi/coco2014'  # folder with data files saved by create_input_files.py\n",
    "data_name = 'coco_5_cap_per_img_5_min_word_freq'  # base name shared by data files\n",
    "checkpoint = './checkpoint_coco_5_cap_per_img_5_min_word_freq.pth.tar'\n",
    "word_map_file = '/data3/zhangweiyi/coco2014/WORDMAP_coco_5_cap_per_img_5_min_word_freq.json'\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "cudnn.benchmark = True\n",
    "beam_size = 1\n",
    "attention = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model\n",
    "checkpoint = torch.load(checkpoint)\n",
    "encoder = checkpoint['encoder']\n",
    "encoder = encoder.to(device)\n",
    "encoder.eval()\n",
    "decoder = checkpoint['decoder']\n",
    "decoder = decoder.to(device)\n",
    "decoder.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load word map (word2ix)\n",
    "with open(word_map_file, 'r') as j:\n",
    "    word_map = json.load(j)\n",
    "rev_word_map = {v: k for k, v in word_map.items()}\n",
    "vocab_size = len(word_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalization transform\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DataLoader\n",
    "loader = torch.utils.data.DataLoader(\n",
    "    CaptionDataset(data_folder, data_name, 'TEST', transform=transforms.Compose([normalize])),\n",
    "    batch_size=1, shuffle=False, num_workers=1, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lists to store references (true captions), and hypothesis (prediction) for each image\n",
    "# If for n images, we have n hypotheses, and references a, b, c... for each image, we need -\n",
    "# references = [[ref1a, ref1b, ref1c], [ref2a, ref2b], ...], hypotheses = [hyp1, hyp2, ...]\n",
    "references = list()\n",
    "hypotheses = list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each image\n",
    "for i, (image, caps, caplens, allcaps) in enumerate(tqdm(loader, desc=\"EVALUATING AT BEAM SIZE \" + str(beam_size))):\n",
    "    k = beam_size\n",
    "\n",
    "    # Tensor to store top k previous words at each step; now they're just <start>\n",
    "    k_prev_words = torch.LongTensor([[word_map['<start>']]] * k).to(device)  # (k, 1)\n",
    "\n",
    "    # Tensor to store top k sequences; now they're just <start>\n",
    "    seqs = k_prev_words  # (k, 1)\n",
    "\n",
    "    # Tensor to store top k sequences' scores; now they're just 0\n",
    "    top_k_scores = torch.zeros(k, 1).to(device)  # (k, 1)\n",
    "\n",
    "    # Lists to store completed sequences and scores\n",
    "    complete_seqs = list()\n",
    "    complete_seqs_scores = list()\n",
    "    \n",
    "    # Move to GPU device, if available\n",
    "    image = image.to(device)  # (1, 3, 256, 256)\n",
    "    \n",
    "    # Encode\n",
    "    encoder_out = encoder(image)  # (1, enc_image_size, enc_image_size, encoder_dim)\n",
    "\n",
    "    # Flatten encoding\n",
    "    # We'll treat the problem as having a batch size of k\n",
    "    if attention:\n",
    "        encoder_dim = encoder_out.size(3)\n",
    "        encoder_out = encoder_out.view(1, -1, encoder_dim)  # (1, num_pixels, encoder_dim)\n",
    "        num_pixels = encoder_out.size(1)\n",
    "        encoder_out = encoder_out.expand(k, num_pixels, encoder_dim)  # (k, num_pixels, encoder_dim)\n",
    "    else:\n",
    "        encoder_out = encoder_out.reshape(1, -1)\n",
    "        encoder_dim = encoder_out.size(1)\n",
    "        encoder_out = encoder_out.expand(k, encoder_dim)\n",
    "    \n",
    "    # Start decoding\n",
    "    step = 1\n",
    "    if attention:\n",
    "        mean_encoder_out = encoder_out.mean(dim=1)\n",
    "        h = decoder.init_h(mean_encoder_out)  # (1, decoder_dim)\n",
    "        c = decoder.init_c(mean_encoder_out)\n",
    "    else:\n",
    "        init_input = decoder.bn(decoder.init(encoder_out))\n",
    "        h, c = decoder.decode_step(init_input)  # (batch_size_t, decoder_dim)\n",
    "    \n",
    "    smoth_wrong = False\n",
    "    \n",
    "    # s is a number less than or equal to k, because sequences are removed from this process once they hit <end>\n",
    "    while True:\n",
    "        embeddings = decoder.embedding(k_prev_words).squeeze(1)  # (s, embed_dim)\n",
    "        if attention:\n",
    "            scores, _, h, c = decoder.one_step(embeddings, encoder_out, h, c)\n",
    "        else:\n",
    "            scores, h, c = decoder.one_step(embeddings, h, c)\n",
    "        scores = F.log_softmax(scores, dim=1)\n",
    "        scores = top_k_scores.expand_as(scores) + scores  # (s, vocab_size)\n",
    "        # For the first step, all k points will have the same scores (since same k previous words, h, c)\n",
    "        if step == 1:\n",
    "            top_k_scores, top_k_words = scores[0].topk(k, 0, True, True)  # (s)\n",
    "        else:\n",
    "            # Unroll and find top scores, and their unrolled indices\n",
    "            top_k_scores, top_k_words = scores.view(-1).topk(k, 0, True, True)  # (s)\n",
    "        # Convert unrolled indices to actual indices of scores\n",
    "        prev_word_inds = top_k_words // vocab_size  # (s)\n",
    "        next_word_inds = top_k_words % vocab_size  # (s)\n",
    "        # Add new words to sequences\n",
    "        seqs = torch.cat([seqs[prev_word_inds], next_word_inds.unsqueeze(1)], dim=1)  # (s, step+1)\n",
    "        # Which sequences are incomplete (didn't reach <end>)?\n",
    "        incomplete_inds = [ind for ind, next_word in enumerate(next_word_inds) if\n",
    "                           next_word != word_map['<end>']]\n",
    "        complete_inds = list(set(range(len(next_word_inds))) - set(incomplete_inds))\n",
    "        # Set aside complete sequences\n",
    "        if len(complete_inds) > 0:\n",
    "            complete_seqs.extend(seqs[complete_inds].tolist())\n",
    "            complete_seqs_scores.extend(top_k_scores[complete_inds])\n",
    "        k -= len(complete_inds)  # reduce beam length accordingly\n",
    "        # Proceed with incomplete sequences\n",
    "        if k == 0:\n",
    "            break\n",
    "        seqs = seqs[incomplete_inds]\n",
    "        h = h[prev_word_inds[incomplete_inds]]\n",
    "        c = c[prev_word_inds[incomplete_inds]]\n",
    "        encoder_out = encoder_out[prev_word_inds[incomplete_inds]]\n",
    "        top_k_scores = top_k_scores[incomplete_inds].unsqueeze(1)\n",
    "        k_prev_words = next_word_inds[incomplete_inds].unsqueeze(1)\n",
    "        # Break if things have been going on too long\n",
    "        if step > 50:\n",
    "            smoth_wrong = True\n",
    "            break\n",
    "        step += 1\n",
    "    if not smoth_wrong:\n",
    "        i = complete_seqs_scores.index(max(complete_seqs_scores))\n",
    "        seq = complete_seqs[i]\n",
    "    else:\n",
    "        seq = seqs[0][:20]\n",
    "    # References\n",
    "    img_caps = allcaps[0].tolist()\n",
    "    img_captions = list(\n",
    "        map(lambda c: [w for w in c if w not in {word_map['<start>'], word_map['<end>'], word_map['<pad>']}],\n",
    "            img_caps))  # remove <start> and pads\n",
    "    references.append(img_captions)\n",
    "    # Hypotheses\n",
    "    hypotheses.append([w for w in seq if w not in {word_map['<start>'], word_map['<end>'], word_map['<pad>']}])\n",
    "    assert len(references) == len(hypotheses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate BLEU-4 scores\n",
    "bleu4 = corpus_bleu(references, hypotheses)\n",
    "print(bleu4)"
   ]
  },
  {
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}