train_program.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "\n",
    "import time\n",
    "import os\n",
    "import copy\n",
    "import json\n",
    "import nltk\n",
    "\n",
    "from tqdm import *\n",
    "from sklearn.metrics import average_precision_score, precision_recall_curve\n",
    "\n",
    "from torch.utils.data import Dataset\n",
    "import skimage\n",
    "from PIL import Image\n",
    "from copy import deepcopy\n",
    "from sklearn.utils.fixes import signature\n",
    "import time\n",
    "import pickle\n",
    "\n",
    "from preprocess import tokenize\n",
    "import programs\n",
    "from modify_program import eliminate_obj_id\n",
    "from utils import *\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.ion()\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def modify_program(questions):\n",
    "    for q in tqdm(questions):\n",
    "        for k in ['equivalent', 'entailed', 'isBalanced', 'groups', 'semanticStr', 'annotations', 'types', 'fullAnswer']:\n",
    "            if k in q:\n",
    "                del q[k]\n",
    "\n",
    "        program = q['semantic']\n",
    "        program = eliminate_obj_id(program)\n",
    "        q['semantic'] = program\n",
    "\n",
    "def program_to_seq(program, mode):\n",
    "    if mode == 'prefix':\n",
    "        program = programs.list_to_prefix(program)\n",
    "    elif mode == 'postfix':\n",
    "        program = programs.list_to_postfix(program)\n",
    "    return program"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GQA(Dataset):\n",
    "    def __init__(self, split):\n",
    "        super(GQA, self).__init__()\n",
    "        \n",
    "        self.word2vocab_id = dict(zip(question_vocab, range(len(question_vocab))))\n",
    "        self.operation2id = dict(zip(operation_vocab, range(len(operation_vocab))))\n",
    "        self.argument2id = dict(zip(argument_vocab, range(len(argument_vocab))))\n",
    "        \n",
    "        dataroot = '/home/qing/Desktop/Datasets/GQA/'\n",
    "        dataset = json.load(open(os.path.join(dataroot, \"questions1.2/%s_questions.json\"%split)))\n",
    "        for k, v in dataset.items():\n",
    "            v['qid'] = k\n",
    "        dataset = list(dataset.values())\n",
    "        modify_program(dataset)\n",
    "        \n",
    "        max_length = 25\n",
    "        max_op = 9\n",
    "        max_arg = 4\n",
    "        for sample in tqdm(dataset):\n",
    "            # 0 -> UNK, 1 -> START, 2-> END\n",
    "            # question\n",
    "            question = sample['question']\n",
    "            tokens = tokenize(question)\n",
    "            tokens = [self.word2vocab_id[x] for x in tokens if x in self.word2vocab_id]\n",
    "            tokens = tokens[:max_length]\n",
    "            padding = []\n",
    "            if len(tokens) < max_length:\n",
    "                padding = [0] * (max_length - len(tokens)) \n",
    "            tokens = padding + [1] + tokens + [2] # Note here we pad in front of the sentence\n",
    "            sample['q_token'] = tokens\n",
    "            \n",
    "            # program\n",
    "            program = program_to_seq(sample['semantic'], 'prefix')\n",
    "            #sample['program_len'] = np.maximum(max_op, len(program)) # program_len does not count 'START' 'END'\n",
    "            # operation\n",
    "            operations = [self.operation2id[x['operation']] for x in program]\n",
    "            operations = operations[:max_op]\n",
    "            padding = []\n",
    "            if len(operations) < max_op:\n",
    "                padding = [0] * (max_op - len(operations))\n",
    "            operations = [1] + operations + [2] + padding\n",
    "            sample['operations'] = operations\n",
    "            \n",
    "            # argument\n",
    "            arguments = []\n",
    "            for op in program:\n",
    "                arg = [self.argument2id[x] for x in op['argument']]\n",
    "                if len(arg) < max_arg:\n",
    "                    padding = [0] * (max_arg - len(arg)) \n",
    "                    arg = arg + padding\n",
    "                arguments.append(arg)\n",
    "                    \n",
    "            padding = []\n",
    "            if len(arguments) < max_op:\n",
    "                padding = [[0] * max_arg] * (max_op - len(arguments))\n",
    "            arguments = [[0] * max_arg] + arguments + [[0] * max_arg] + padding\n",
    "            sample['arguments'] = arguments\n",
    "            \n",
    "            del sample['question'], sample['semantic']\n",
    "            \n",
    "            \n",
    "        self.dataset = dataset\n",
    "        \n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        entry = self.dataset[index]\n",
    "        question = torch.LongTensor(np.array(entry['q_token']))\n",
    "        operations = torch.LongTensor(np.array(entry['operations']))\n",
    "        arguments = torch.LongTensor(np.array(entry['arguments']))\n",
    "        return question, operations, arguments\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 132062/132062 [00:00<00:00, 148025.90it/s]\n",
      "100%|██████████| 132062/132062 [00:03<00:00, 39127.78it/s]\n",
      "100%|██████████| 943000/943000 [00:06<00:00, 136158.65it/s]\n",
      "100%|██████████| 943000/943000 [00:25<00:00, 36931.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'val_balanced': 132062, 'train_balanced': 943000}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\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)\n",
    "\n",
    "question_vocab = json.load(open('question_vocab.json'))\n",
    "operation_vocab = json.load(open('operation_vocab.json'))\n",
    "argument_vocab = json.load(open('argument_vocab.json'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, in_dim, num_hid, nlayers = 1, dropout = 0., rnn_type='GRU'):\n",
    "        \"\"\"Module for decoder\n",
    "        \"\"\"\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        assert rnn_type == 'LSTM' or rnn_type == 'GRU'\n",
    "        rnn_cls = nn.LSTM if rnn_type == 'LSTM' else nn.GRU\n",
    "\n",
    "        self.rnn = rnn_cls(\n",
    "            in_dim, num_hid, nlayers,\n",
    "            dropout=dropout,\n",
    "            batch_first=True)\n",
    "\n",
    "        self.in_dim = in_dim\n",
    "        self.num_hid = num_hid\n",
    "        self.nlayers = nlayers\n",
    "        self.rnn_type = rnn_type\n",
    "\n",
    "    def init_hidden(self, batch):\n",
    "        # just to get the type of tensor\n",
    "        weight = next(self.parameters()).data\n",
    "        hid_shape = (self.nlayers, batch, self.num_hid)\n",
    "        if self.rnn_type == 'LSTM':\n",
    "            return (Variable(weight.new(*hid_shape).zero_()),\n",
    "                    Variable(weight.new(*hid_shape).zero_()))\n",
    "        else:\n",
    "            return Variable(weight.new(*hid_shape).zero_())\n",
    "\n",
    "    def forward(self, input, hidden):\n",
    "        # input: [batch, sequence, in_dim]\n",
    "        self.rnn.flatten_parameters()\n",
    "        output, hidden = self.rnn(input, hidden)\n",
    "        return output, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "from language_model import WordEmbedding, QuestionEmbedding\n",
    "from torch.nn.utils.weight_norm import weight_norm\n",
    "\n",
    "class ProgramGenerator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ProgramGenerator, self).__init__()\n",
    "        self.w_emb = WordEmbedding(len(question_vocab), 300)\n",
    "        self.w_emb.init_embedding('question_word_embedding_glove_init.npy')\n",
    "        self.op_emb = WordEmbedding(len(operation_vocab), 300)\n",
    "        self.arg_emb = WordEmbedding(len(argument_vocab), 300)\n",
    "        \n",
    "        self.q_encoder = QuestionEmbedding(300, 512)\n",
    "        self.p_decoder = DecoderRNN(300*5, 512) # 5 = 1 + 4, 1 operation, 4 arguments\n",
    "        self.op_predictor = weight_norm(nn.Linear(512, len(operation_vocab)), dim=None)\n",
    "        self.arg_predictor = weight_norm(nn.Linear(512, 4 * len(argument_vocab)), dim=None)\n",
    "\n",
    "\n",
    "    def forward(self, question, operations, arguments):\n",
    "        \"\"\"Forward\n",
    "        return: logits, not probs\n",
    "        \"\"\"\n",
    "        \n",
    "        batch_size = question.size(0)\n",
    "        question_len = question.size(1)\n",
    "        program_len = operations.size(1)\n",
    "        \n",
    "        w_emb = self.w_emb(question)\n",
    "        q_enc = self.q_encoder(w_emb) # [batch, q_dim]\n",
    "        \n",
    "        \n",
    "        operation_emb = self.op_emb(operations)\n",
    "        argument_emb = self.arg_emb(arguments)\n",
    "        \n",
    "        argument_emb = argument_emb.view(batch_size, program_len, -1)\n",
    "        op_arg_emb = torch.cat((operation_emb, argument_emb), -1)\n",
    "        \n",
    "        if self.training:\n",
    "            # Teacher forcing: Feed the target as the next input\n",
    "            p_dec, _ = self.p_decoder(op_arg_emb, q_enc.view(1, batch_size, -1))\n",
    "            op_logit = self.op_predictor(p_dec) # [batch, program_len, operation_vocab]\n",
    "            arg_logit = self.arg_predictor(p_dec) \n",
    "            arg_logit = arg_logit.view(batch_size, program_len, 4, -1) # [batch, program_len, 4, arg_vocab]\n",
    "        \n",
    "        else:\n",
    "            # Without teacher forcing: use its own predictions as the next input\n",
    "            op_logit = []\n",
    "            arg_logit = []\n",
    "            decoder_hidden = q_enc.view(1, batch_size, -1)\n",
    "            decoder_input = op_arg_emb[:, 0].view(batch_size, 1, -1)\n",
    "            for di in range(program_len):\n",
    "                decoder_output, decoder_hidden = self.p_decoder(\n",
    "                    decoder_input, decoder_hidden)\n",
    "\n",
    "                one_op_logit = self.op_predictor(decoder_output)\n",
    "                one_arg_logit = self.arg_predictor(decoder_output)\n",
    "                one_arg_logit = one_arg_logit.view(batch_size, 4, -1)\n",
    "\n",
    "                _, op_pred = one_op_logit.topk(1, dim=-1)\n",
    "                _, arg_pred = one_arg_logit.topk(1, dim=-1)\n",
    "                op_pred = op_pred.view(batch_size)\n",
    "                arg_pred = arg_pred.view(batch_size, -1)\n",
    "                operation_emb = self.op_emb(op_pred)\n",
    "                argument_emb = self.arg_emb(arg_pred)\n",
    "                argument_emb = argument_emb.view(batch_size, -1)\n",
    "                \n",
    "                op_arg_emb = torch.cat((operation_emb, argument_emb), -1)\n",
    "                decoder_input = op_arg_emb.view(batch_size, 1, -1).detach()  # detach from history as input\n",
    "                \n",
    "                op_logit.append(one_op_logit)\n",
    "                arg_logit.append(one_arg_logit)\n",
    "            op_logit = torch.stack(op_logit, dim=1)\n",
    "            arg_logit = torch.stack(arg_logit, dim=1)\n",
    "        \n",
    "        return op_logit, arg_logit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "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",
    "    mask = gt.data != 0 # pad with 0\n",
    "    prob = prob[mask, :]\n",
    "    gt = gt[mask]\n",
    "    loss = F.cross_entropy(prob, gt)\n",
    "    \n",
    "    pred = prob.max(dim=1)[1]\n",
    "    acc = (pred == gt).float().mean()\n",
    "    return loss, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, dataloader):\n",
    "    model.eval() \n",
    "    score_all = []\n",
    "    label_all = []\n",
    "\n",
    "    # Iterate over data.\n",
    "    op_correct = 0\n",
    "    arg_correct = 0\n",
    "    total_count = 0\n",
    "    for question, operations, arguments in dataloader:\n",
    "        question = question.to(device)\n",
    "        operations = operations.to(device)\n",
    "        arguments = arguments.to(device)\n",
    "        op_logit, arg_logit = model(question, operations[:, :-1], arguments[:, :-1])\n",
    "        op_loss, op_acc = compute_loss_acc(\n",
    "            op_logit.view(-1, len(operation_vocab)),\n",
    "            operations[:, 1:].contiguous().view(-1))\n",
    "        arg_loss, arg_acc = compute_loss_acc(\n",
    "            arg_logit.view(-1, len(argument_vocab)),\n",
    "            arguments[:, 1:].contiguous().view(-1))\n",
    "        \n",
    "        batch_size = question.size(0)\n",
    "        op_correct += op_acc * batch_size\n",
    "        arg_correct += arg_acc * batch_size\n",
    "        total_count += batch_size\n",
    "    op_acc = op_correct / total_count\n",
    "    arg_acc = arg_correct / total_count\n",
    "    \n",
    "    return op_acc, arg_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "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=128,\n",
    "                                             shuffle=True, num_workers=4) for x in train_splits}\n",
    "    \n",
    "    eval_dataloaders = {x: torch.utils.data.DataLoader(datasets[x], batch_size=128,\n",
    "                                         shuffle=False, num_workers=4) 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",
    "        op_acc, arg_acc = evaluate_model(model, dataloaders[eval_split])\n",
    "        print('(op_acc={1:.2f}, arg_acc={2:.2f}) {0}'.format(eval_split, 100*op_acc, 100*arg_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 question, operations, arguments in dataloaders[train_split]:\n",
    "                model.train()  # Set model to training mode\n",
    "                question = question.to(device)\n",
    "                operations = operations.to(device)\n",
    "                arguments = arguments.to(device)\n",
    "                op_logit, arg_logit = model(question, operations[:, :-1], arguments[:, :-1])\n",
    "                \n",
    "                op_loss, op_acc = compute_loss_acc(\n",
    "                    op_logit.view(-1, len(operation_vocab)),\n",
    "                    operations[:, 1:].contiguous().view(-1))\n",
    "                arg_loss, arg_acc = compute_loss_acc(\n",
    "                    arg_logit.view(-1, len(argument_vocab)),\n",
    "                    arguments[:, 1:].contiguous().view(-1))\n",
    "                \n",
    "                loss = op_loss + arg_loss\n",
    "\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "\n",
    "        print(op_acc.item(), arg_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",
    "                op_acc, arg_acc = evaluate_model(model, dataloaders[eval_split])\n",
    "                print('(op_acc={1:.2f}, arg_acc={2:.2f}) {0}'.format(eval_split, 100*op_acc, 100*arg_acc))\n",
    "            acc = op_acc + arg_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",
    "        ap, label, score = evaluate_model(model, dataloaders[eval_split])\n",
    "        print('(AP={1:.2f}) {0}'.format(eval_split, 100*ap))\n",
    "    # deep copy the model\n",
    "    if ap > best_ap:\n",
    "        best_ap = ap\n",
    "        best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    print()\n",
    "    #########################################\n",
    "\n",
    "    print('Best val AP: {:2f}'.format(100*best_ap))\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(op_acc=0.05, arg_acc=0.00) val_balanced\n",
      "\n",
      "Epoch 0/19\n",
      "----------\n",
      "0.9897959232330322 0.7685949802398682 1.0681862831115723\n",
      "(op_acc=66.69, arg_acc=45.78) val_balanced\n",
      "Epoch time: 2m 32s\n",
      "\n",
      "Epoch 1/19\n",
      "----------\n",
      "1.0 0.8869564533233643 0.455322802066803\n",
      "(op_acc=67.01, arg_acc=55.52) val_balanced\n",
      "Epoch time: 2m 33s\n",
      "\n",
      "Epoch 2/19\n",
      "----------\n",
      "1.0 0.8691588640213013 0.38715264201164246\n",
      "(op_acc=60.86, arg_acc=58.14) val_balanced\n",
      "Epoch time: 2m 37s\n",
      "\n",
      "Epoch 3/19\n",
      "----------\n",
      "1.0 0.9711538553237915 0.09018203616142273\n",
      "(op_acc=59.64, arg_acc=58.89) val_balanced\n",
      "Epoch time: 2m 40s\n",
      "\n",
      "Epoch 4/19\n",
      "----------\n",
      "0.9999999403953552 0.9385964870452881 0.1323327273130417\n",
      "(op_acc=58.15, arg_acc=58.92) val_balanced\n",
      "Epoch time: 2m 40s\n",
      "\n",
      "Epoch 5/19\n",
      "----------\n",
      "0.9909909963607788 0.9752065539360046 0.13251852989196777\n",
      "(op_acc=58.41, arg_acc=58.56) val_balanced\n",
      "Epoch time: 2m 36s\n",
      "\n",
      "Epoch 6/19\n",
      "----------\n",
      "0.9999999403953552 0.9523810148239136 0.08971542119979858\n",
      "(op_acc=61.00, arg_acc=59.92) val_balanced\n",
      "Epoch time: 2m 39s\n",
      "\n",
      "Epoch 7/19\n",
      "----------\n",
      "0.9999999403953552 0.9479166865348816 0.1744961440563202\n",
      "(op_acc=58.62, arg_acc=61.19) val_balanced\n",
      "Epoch time: 2m 36s\n",
      "\n",
      "Epoch 8/19\n",
      "----------\n",
      "1.0 0.9357797503471375 0.1583596169948578\n",
      "(op_acc=57.94, arg_acc=60.92) val_balanced\n",
      "Epoch time: 2m 36s\n",
      "\n",
      "Epoch 9/19\n",
      "----------\n",
      "0.9999999403953552 0.8999999761581421 0.16963811218738556\n",
      "(op_acc=58.43, arg_acc=61.83) val_balanced\n",
      "Epoch time: 2m 32s\n",
      "\n",
      "Epoch 10/19\n",
      "----------\n",
      "1.0 0.9523810148239136 0.14111706614494324\n",
      "(op_acc=58.45, arg_acc=61.82) val_balanced\n",
      "Epoch time: 2m 41s\n",
      "\n",
      "Epoch 11/19\n",
      "----------\n",
      "0.9999999403953552 0.9729729890823364 0.07740219682455063\n",
      "(op_acc=58.65, arg_acc=63.07) val_balanced\n",
      "Epoch time: 2m 36s\n",
      "\n",
      "Epoch 12/19\n",
      "----------\n",
      "1.0 0.9734513163566589 0.05603726953268051\n",
      "(op_acc=59.29, arg_acc=62.88) val_balanced\n",
      "Epoch time: 2m 32s\n",
      "\n",
      "Epoch 13/19\n",
      "----------\n",
      "1.0 0.9739130139350891 0.08276142925024033\n",
      "(op_acc=59.31, arg_acc=64.06) val_balanced\n",
      "Epoch time: 2m 34s\n",
      "\n",
      "Epoch 14/19\n",
      "----------\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-82-ecab2f2d3132>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0meval_splits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m'val_balanced'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mProgramGenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mtrain_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_epochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m20\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_splits\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrain_splits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0meval_splits\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0meval_splits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn_epochs_per_eval\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-81-99083da0e255>\u001b[0m in \u001b[0;36mtrain_model\u001b[0;34m(model, num_epochs, train_splits, eval_splits, n_epochs_per_eval)\u001b[0m\n\u001b[1;32m     73\u001b[0m                 op_loss, op_acc = compute_loss_acc(\n\u001b[1;32m     74\u001b[0m                     \u001b[0mop_logit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moperation_vocab\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 75\u001b[0;31m                     operations[:, 1:].contiguous().view(-1))\n\u001b[0m\u001b[1;32m     76\u001b[0m                 arg_loss, arg_acc = compute_loss_acc(\n\u001b[1;32m     77\u001b[0m                     \u001b[0marg_logit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margument_vocab\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-75-f3bdbbd8c977>\u001b[0m in \u001b[0;36mcompute_loss_acc\u001b[0;34m(prob, gt)\u001b[0m\n\u001b[1;32m      6\u001b[0m     \"\"\"\n\u001b[1;32m      7\u001b[0m     \u001b[0mmask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;31m# pad with 0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m     \u001b[0mprob\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprob\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m     \u001b[0mgt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgt\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcross_entropy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprob\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train_splits = ['train_balanced']\n",
    "eval_splits = ['val_balanced']\n",
    "model = ProgramGenerator().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": []
  },
  {
   "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
}