train_execution.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import math\n",
    "import numpy as np\n",
    "import copy\n",
    "from tqdm import *\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.ion()\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models\n",
    "from torch.nn.utils.weight_norm import weight_norm\n",
    "from torch import optim\n",
    "from torch.optim import lr_scheduler\n",
    "\n",
    "from language_model import WordEmbedding, QuestionEmbedding\n",
    "from fc import FCNet\n",
    "from utils import *\n",
    "from modify_program import *\n",
    "from dataset import *\n",
    "\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "#device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_answer_prob(yes_prob):\n",
    "    yes_id = answer2id['yes']\n",
    "    no_id = answer2id['no']\n",
    "    prob = torch.zeros((1, len(answer_vocab))).to(device)\n",
    "    prob[0, yes_id] = yes_prob\n",
    "    prob[0, no_id] = 1. - yes_prob\n",
    "    return prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class And(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(And, self).__init__()\n",
    "    \n",
    "    def forward(self, v, p0, p1):\n",
    "        yes_id = answer2id['yes']\n",
    "        p0 = p0[0, yes_id]\n",
    "        p1 = p1[0, yes_id]\n",
    "        yes_prob = p0 * p1\n",
    "        prob = make_answer_prob(yes_prob)\n",
    "        return prob\n",
    "    \n",
    "class Or(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Or, self).__init__()\n",
    "    \n",
    "    def forward(self, v, p0, p1):\n",
    "        yes_id = answer2id['yes']\n",
    "        p0 = p0[0, yes_id]\n",
    "        p1 = p1[0, yes_id]\n",
    "        yes_prob = p0 + p1 - p0 * p1\n",
    "        prob = make_answer_prob(yes_prob)\n",
    "        return prob\n",
    "    \n",
    "class AttentionAnd(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(AttentionAnd, self).__init__()\n",
    "        \n",
    "    def forward(self, v, a1, a2):\n",
    "        return torch.min(a1, a2)\n",
    "    \n",
    "class AttentionNot(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(AttentionNot, self).__init__()\n",
    "        \n",
    "    def forward(self, v, a):\n",
    "        return 1. - a\n",
    "    \n",
    "class AttentionOr(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(AttentionOr, self).__init__()\n",
    "        \n",
    "    def forward(self, v, a1, a2):\n",
    "        return torch.max(a1, a2)\n",
    "\n",
    "class Exist(nn.Module):\n",
    "    def __init__(self, att_size=100):\n",
    "        super(Exist, self).__init__()\n",
    "        self.linear = weight_norm(nn.Linear(att_size, 1), dim=None)\n",
    "    \n",
    "    def forward(self, v, att, arg):\n",
    "        batch = att.size(0)\n",
    "        att = att.view(batch, -1)\n",
    "        logits = self.linear(att)\n",
    "        logits = torch.sigmoid(logits)\n",
    "        prob = make_answer_prob(logits)\n",
    "        return prob\n",
    "    \n",
    "class Choose(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Choose, self).__init__()\n",
    "    \n",
    "    def forward(self, v, p0, p1, choice0, choice1):\n",
    "        yes_id = answer2id['yes']\n",
    "        p0 = p0[0, yes_id]\n",
    "        p1 = p1[0, yes_id]\n",
    "        p0 /= p0 + p1\n",
    "        p1 /= p0 + p1\n",
    "        prob = torch.zeros((1, len(answer_vocab),)).to(device)\n",
    "        \n",
    "        mapping = {'to the left of': 'left',\n",
    "                    'to the right of': 'right',\n",
    "                    'in front of': 'front', \n",
    "                    'standing in front of': 'front'}\n",
    "        choice0 = mapping.get(choice0, choice0)\n",
    "        choice1 = mapping.get(choice1, choice1)\n",
    "            \n",
    "        prob[0, answer2id.get(choice0, 0)] = p0 # 0 -> UNK\n",
    "        prob[0, answer2id.get(choice1, 0)] = p1\n",
    "        return prob\n",
    "    \n",
    "\n",
    "class Select(nn.Module):\n",
    "    def __init__(self, v_dim=2048, t_dim=512, num_hid=512, dropout=0.):\n",
    "        super(Select, self).__init__()\n",
    "\n",
    "        self.v_proj = FCNet([v_dim, num_hid])\n",
    "        self.t_proj = FCNet([t_dim, num_hid])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear = weight_norm(nn.Linear(num_hid, 1), dim=None)\n",
    "\n",
    "    def forward(self, v, t):\n",
    "        \"\"\"\n",
    "        v: [batch, k, v_dim]\n",
    "        t: [batch, t_dim]\n",
    "        \"\"\"\n",
    "        batch, k, _ = v.size()\n",
    "        v_proj = self.v_proj(v)\n",
    "        t_proj = self.t_proj(t).unsqueeze(1).repeat(1, k, 1)\n",
    "        joint_repr = v_proj * t_proj\n",
    "        joint_repr = self.dropout(joint_repr)\n",
    "        logits = self.linear(joint_repr)\n",
    "        #w = F.softmax(logits, 1)\n",
    "        w = torch.sigmoid(logits)\n",
    "        return w\n",
    "        \n",
    "class Relocate(nn.Module):\n",
    "    def __init__(self, v_dim=2048, t_dim=512, num_hid=512, dropout=0.):\n",
    "        super(Relocate, self).__init__()\n",
    "\n",
    "        self.v_proj = FCNet([v_dim, num_hid])\n",
    "        self.t_proj = FCNet([t_dim, num_hid])\n",
    "        self.av_proj = FCNet([v_dim, num_hid])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear = weight_norm(nn.Linear(num_hid, 1), dim=None)\n",
    "\n",
    "    def forward(self, v, a, t, so):\n",
    "        \"\"\"\n",
    "        v: [batch, k, v_dim] vis\n",
    "        a: [batch, k] attention\n",
    "        t: [batch, t_dim] txt\n",
    "        so: [batch, t_dim] subject, object\n",
    "        \"\"\"\n",
    "        batch, k, _ = v.size()\n",
    "\n",
    "        v_proj = self.v_proj(v)\n",
    "        t_proj = self.t_proj(t).unsqueeze(1).repeat(1, k, 1)\n",
    "        \n",
    "        av = (a * v).sum(1)\n",
    "        av_proj = self.av_proj(av).unsqueeze(1).repeat(1, k, 1)\n",
    "        \n",
    "        joint_repr = v_proj * t_proj * av_proj\n",
    "        joint_repr = self.dropout(joint_repr)\n",
    "        logits = self.linear(joint_repr)\n",
    "        #w = F.softmax(logits, 1)\n",
    "        w = torch.sigmoid(logits)\n",
    "        return w\n",
    "        \n",
    "class Compare(nn.Module):\n",
    "    def __init__(self, v_dim=2048, t_dim=512, num_hid=512, dropout=0.):\n",
    "        super(Compare, self).__init__()\n",
    "\n",
    "        self.v_proj = FCNet([v_dim, num_hid])\n",
    "        self.t_proj = FCNet([t_dim, num_hid])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear = weight_norm(nn.Linear(num_hid, 1), dim=None)\n",
    "\n",
    "    def forward(self, v, a, t1, t2):\n",
    "        \"\"\"\n",
    "        v: [batch, k, v_dim]\n",
    "        a: [batch, k]\n",
    "        t1: [batch, t_dim] e.g., different, same, ...\n",
    "        t2: [batch, t_dim] e.g., color, type, material, ...\n",
    "        \"\"\"\n",
    "        batch, k, _ = v.size()\n",
    "        v = (a * v).sum(1)\n",
    "        v = self.v_proj(v)\n",
    "        t1 = self.t_proj(t1)\n",
    "        t2 = self.t_proj(t2)\n",
    "        joint_repr = v * t1 * t2\n",
    "        joint_repr = self.dropout(joint_repr)\n",
    "        logits = self.linear(joint_repr)\n",
    "        logits = torch.sigmoid(logits)\n",
    "        prob = make_answer_prob(logits)\n",
    "        return prob\n",
    "    \n",
    "class Common(nn.Module):\n",
    "    def __init__(self, v_dim=2048, t_dim=512, num_hid=512, dropout=0.):\n",
    "        super(Common, self).__init__()\n",
    "\n",
    "        self.v_proj = FCNet([v_dim, num_hid])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear = weight_norm(nn.Linear(num_hid, len(answer_vocab)), dim=None)\n",
    "\n",
    "    def forward(self, v, a1, a2):\n",
    "        \"\"\"\n",
    "        v: [batch, k, v_dim]\n",
    "        a1: [batch, k]\n",
    "        a2: [batch, k]\n",
    "        \"\"\"\n",
    "        batch, k, _ = v.size()\n",
    "        v = self.v_proj(v)\n",
    "        av1 = (a1 * v).sum(1)\n",
    "        av2 = (a2 * v).sum(1)\n",
    "        joint_repr = av1 * av2\n",
    "        joint_repr = self.dropout(joint_repr)\n",
    "        logits = self.linear(joint_repr)\n",
    "        prob = F.softmax(logits, 1)\n",
    "        return prob\n",
    "    \n",
    "class Query(nn.Module):\n",
    "    def __init__(self, v_dim=2048, t_dim=512, num_hid=512, dropout=0.):\n",
    "        super(Query, self).__init__()\n",
    "\n",
    "        self.v_proj = FCNet([v_dim, num_hid])\n",
    "        self.t_proj = FCNet([t_dim, num_hid])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear = weight_norm(nn.Linear(num_hid, len(answer_vocab)), dim=None)\n",
    "\n",
    "    def forward(self, v, a, t):\n",
    "        \"\"\"\n",
    "        v: [batch, k, v_dim] \n",
    "        a: [batch, k]\n",
    "        t: [batch, t_dim]\n",
    "        \"\"\"\n",
    "        batch, k, _ = v.size()\n",
    "        v = self.v_proj(v)\n",
    "        v = (a * v).sum(1)\n",
    "        t = self.t_proj(t)\n",
    "        joint_repr = v * t\n",
    "        joint_repr = self.dropout(joint_repr)\n",
    "        logits = self.linear(joint_repr)\n",
    "        prob = F.softmax(logits, 1)\n",
    "        return prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModuleNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ModuleNet, self).__init__()\n",
    "        self.function_modules = {}\n",
    "        \n",
    "        Modules = ['And', 'Or', 'AttentionAnd', 'AttentionNot', 'AttentionOr', 'Exist', 'Choose', 'Compare', \n",
    "                   'Common', 'Query', 'Select', 'Relocate']\n",
    "        # Initialize all modules\n",
    "        for module in Modules:\n",
    "            func_net = eval(module)()\n",
    "            self.add_module(module, func_net)\n",
    "            self.function_modules[module.lower()] = func_net\n",
    "            \n",
    "        self.arg_emb = WordEmbedding(len(argument_vocab), 512)\n",
    "    \n",
    "    def forward(self, img_feats, program):\n",
    "        N = img_feats.size(0)\n",
    "        final_module_outputs = []\n",
    "        for i in range(N):\n",
    "            module_outputs = []\n",
    "            for j, f in enumerate(program[i]):\n",
    "                #print(f)\n",
    "                module = self.function_modules[f['operation']]\n",
    "                module_inputs = [img_feats[i:i+1]]\n",
    "                module_inputs.extend([module_outputs[dep] for dep in f['dependencies']])\n",
    "                \n",
    "                if f['operation'] == 'choose': # for choose operation, argument, not the argument embedding\n",
    "                    assert len(f['argument']) == 2\n",
    "                    module_inputs.extend(f['argument'])\n",
    "                else:                    \n",
    "                    module_inputs.extend([self.arg_emb(torch.LongTensor([argument2id[arg]]).to(device))\n",
    "                                          for arg in f['argument']])\n",
    "                    \n",
    "                module_outputs.append(module(*module_inputs))\n",
    "            final_module_outputs.append(module_outputs[-1])\n",
    "        \n",
    "        final_module_outputs = torch.cat(final_module_outputs, 0)\n",
    "        return final_module_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'val_balanced': 132062, 'train_balanced': 943000}\n"
     ]
    }
   ],
   "source": [
    "splits = ['val_balanced', 'train_balanced']\n",
    "datasets = {}\n",
    "datasets.update({x: GQA(x) for x in splits})\n",
    "dataset_sizes = {x: len(datasets[x]) for x in splits}\n",
    "print(dataset_sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss_acc(prob, gt):\n",
    "    \"\"\"\n",
    "    Inputs:\n",
    "    - y_pred: Variable of shape (N, V_out)\n",
    "    - y: LongTensor Variable of shape (N,)\n",
    "    \"\"\"\n",
    "    loss = F.nll_loss(torch.log(prob+1e-7), gt)\n",
    "    \n",
    "    pred = prob.max(dim=1)[1]\n",
    "    acc = (pred == gt).float().mean()\n",
    "    return loss, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, dataloader):\n",
    "    model.eval() \n",
    "    # Iterate over data.\n",
    "    correct = 0\n",
    "    total_count = 0\n",
    "    for img_feat, program, answer in tqdm(dataloader):\n",
    "        img_feat = img_feat.to(device)\n",
    "        answer = answer.to(device)\n",
    "        prob = model(img_feat, program)\n",
    "        loss, acc = compute_loss_acc(prob, answer)\n",
    "        batch_size = img_feat.size(0)\n",
    "        correct += acc * batch_size\n",
    "        total_count += batch_size\n",
    "    acc = correct / total_count\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, num_epochs=5, train_splits=['train'], \n",
    "                eval_splits=['val'], n_epochs_per_eval = 1):\n",
    "    optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "    # Decay LR by a factor of 0.1 every 100 epochs\n",
    "    scheduler = lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.1)\n",
    "    \n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "    \n",
    "    train_dataloaders = {x: torch.utils.data.DataLoader(datasets[x], batch_size=32,\n",
    "                         shuffle=True, num_workers=4, collate_fn=GQA_collate) for x in train_splits}\n",
    "    \n",
    "    eval_dataloaders = {x: torch.utils.data.DataLoader(datasets[x], batch_size=32,\n",
    "                         shuffle=False, num_workers=4, collate_fn=GQA_collate) for x in eval_splits}\n",
    "    \n",
    "    dataloaders = {}\n",
    "    dataloaders.update(train_dataloaders)\n",
    "    dataloaders.update(eval_dataloaders)\n",
    "    \n",
    "    ###########evaluate init model###########\n",
    "    for eval_split in eval_splits:\n",
    "        acc = evaluate_model(model, dataloaders[eval_split])\n",
    "        print('(acc={1:.2f}) {0}'.format(eval_split, 100*acc))\n",
    "    print()\n",
    "    #########################################\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        since = time.time()\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "        \n",
    "        scheduler.step()\n",
    "\n",
    "        # Iterate over data.\n",
    "        for train_split in train_splits:\n",
    "            for img_feat, program, answer in tqdm(dataloaders[train_split]):\n",
    "                model.train()  # Set model to training mode\n",
    "                img_feat = img_feat.to(device)\n",
    "                answer = answer.to(device)\n",
    "                prob = model(img_feat, program)\n",
    "                loss, acc = compute_loss_acc(prob, answer)\n",
    "\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "\n",
    "        print(acc.item(), loss.item())\n",
    "        # compute average precision\n",
    "        if (epoch+1) % n_epochs_per_eval == 0:\n",
    "            for eval_split in eval_splits:\n",
    "                acc = evaluate_model(model, dataloaders[eval_split])\n",
    "                print('(acc={1:.2f}) {0}'.format(eval_split, 100*acc))\n",
    "            # deep copy the model\n",
    "            if acc > best_acc:\n",
    "                best_acc = acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "                \n",
    "        time_elapsed = time.time() - since\n",
    "        print('Epoch time: {:.0f}m {:.0f}s'.format(\n",
    "            time_elapsed // 60, time_elapsed % 60))\n",
    "        print(flush=True)\n",
    "    \n",
    "    ###########evaluate final model###########\n",
    "    for eval_split in eval_splits:\n",
    "        acc = evaluate_model(model, dataloaders[eval_split])\n",
    "        print('(acc={1:.2f}) {0}'.format(eval_split, 100*acc))\n",
    "    # deep copy the model\n",
    "    if acc > best_acc:\n",
    "        best_acc = acc\n",
    "        best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    #########################################\n",
    "\n",
    "    print('Best val acc: {:2f}'.format(100*best_acc))\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|          | 0/4127 [00:00<?, ?it/s]\u001b[A\n",
      "  0%|          | 1/4127 [00:03<4:11:06,  3.65s/it]\u001b[A\n",
      "  0%|          | 2/4127 [00:03<2:58:10,  2.59s/it]\u001b[A\n",
      "  0%|          | 4/4127 [00:03<2:06:29,  1.84s/it]\u001b[A\n",
      "  0%|          | 5/4127 [00:04<1:34:51,  1.38s/it]\u001b[A\n",
      "  0%|          | 6/4127 [00:04<1:11:10,  1.04s/it]\u001b[A\n",
      "  0%|          | 8/4127 [00:04<51:39,  1.33it/s]  \u001b[A\n",
      "  0%|          | 9/4127 [00:05<45:11,  1.52it/s]\u001b[A\n",
      "  0%|          | 10/4127 [00:05<35:09,  1.95it/s]\u001b[A\n",
      "  0%|          | 12/4127 [00:05<26:19,  2.61it/s]\u001b[A\n",
      "  0%|          | 13/4127 [00:05<28:40,  2.39it/s]\u001b[A\n",
      "  0%|          | 14/4127 [00:06<27:24,  2.50it/s]\u001b[A\n",
      "  0%|          | 16/4127 [00:06<20:51,  3.29it/s]\u001b[A\n",
      "  0%|          | 17/4127 [00:06<18:39,  3.67it/s]\u001b[A\n",
      "  0%|          | 18/4127 [00:07<26:52,  2.55it/s]\u001b[A\n",
      "  0%|          | 20/4127 [00:07<20:35,  3.32it/s]\u001b[A\n",
      "  1%|          | 22/4127 [00:07<18:00,  3.80it/s]\u001b[A\n",
      "  1%|          | 23/4127 [00:07<15:14,  4.49it/s]\u001b[A\n",
      "  1%|          | 24/4127 [00:08<13:35,  5.03it/s]\u001b[A\n",
      "  1%|          | 25/4127 [00:08<11:47,  5.80it/s]\u001b[A\n",
      "  1%|          | 26/4127 [00:08<15:04,  4.53it/s]\u001b[A\n",
      "  1%|          | 28/4127 [00:09<14:55,  4.58it/s]\u001b[A\n",
      "  1%|          | 30/4127 [00:09<14:01,  4.87it/s]\u001b[A\n",
      "  1%|          | 32/4127 [00:09<13:53,  4.91it/s]\u001b[A\n",
      "  1%|          | 34/4127 [00:10<12:31,  5.45it/s]\u001b[A\n",
      "  1%|          | 35/4127 [00:10<10:52,  6.27it/s]\u001b[A\n",
      "  1%|          | 36/4127 [00:10<16:01,  4.25it/s]\u001b[A\n",
      "  1%|          | 38/4127 [00:10<15:09,  4.50it/s]\u001b[A"
     ]
    }
   ],
   "source": [
    "train_splits = ['train_balanced']\n",
    "eval_splits = ['val_balanced']\n",
    "model = ModuleNet().to(device)\n",
    "train_model(model, num_epochs=20, train_splits=train_splits, eval_splits=eval_splits, n_epochs_per_eval = 1)"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}