a2c-train.py

from __future__ import division
import argparse
import torch
import torch.nn as nn
from torch import cuda
import lib
import os
import sys
import pdb

import datetime
import numpy as np
import os.path
from torch.autograd import Variable
import random
import gensim
import time
import pickle
import code_retrieval

# # to deal with version incompatible
# import torch._utils
# try:
#     torch._utils._rebuild_tensor_v2
# except AttributeError:
#     def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
#         tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
#         tensor.requires_grad = requires_grad
#         tensor._backward_hooks = backward_hooks
#         return tensor
#     torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2


def get_opt():
    parser = argparse.ArgumentParser(description='a2c-train.py')
    # Data options
    parser.add_argument('-data', required=True, help='Path to the *-train.pt file from preprocess.py')
    parser.add_argument('-lang', required=True, choices=['sql', 'python'], help='Language {python|sql}')
    parser.add_argument('-data_name', default="", help="Data name, such as toy")
    parser.add_argument('-save_dir', required=True, help='Directory to save models')
    parser.add_argument("-load_from", help="Path to load a pretrained model.")
    parser.add_argument("-show_str", required=True, help="string of arguments for saving models.")
    parser.add_argument('-load_embedding_from', required=False, help='Path to load the embedding.')

    # Model options
    parser.add_argument('-layers', type=int, default=1, help='Number of layers in the LSTM encoder/decoder')
    parser.add_argument('-rnn_size', type=int, default=512, help='Size of LSTM hidden states')
    parser.add_argument('-word_vec_size', type=int, default=512, help='Word embedding sizes')
    parser.add_argument('-input_feed', type=int, default=1, help="""Feed the context vector at each time step as
                        additional input (via concatenation with the word embeddings) to the decoder.""")
    parser.add_argument('-brnn', action='store_true', help='Use a bidirectional encoder')
    parser.add_argument('-brnn_merge', default='concat', help="""Merge action for the bidirectional hidden states: [concat|sum]""")
    parser.add_argument('-has_attn', type=int, default=1, help="""attn model or not""")
    parser.add_argument('-has_baseline', type=int, default=1, help="baseline model")

    # Optimization options
    parser.add_argument('-batch_size', type=int, default=64, help='Maximum batch size')
    parser.add_argument("-max_generator_batches", type=int, default=128, help="""Split softmax input into small batches for memory efficiency. Higher is faster, but uses more memory.""")

    parser.add_argument("-end_epoch", type=int, default=50, help="Epoch to stop training.")
    parser.add_argument("-start_epoch", type=int, default=1, help="Epoch to start training.")

    parser.add_argument('-param_init', type=float, default=0.1, help="""Parameters are initialized over uniform distribution with support (-param_init, param_init). Use 0 to not use initialization""")
    parser.add_argument('-optim', default='adam', help="Optimization method. [sgd|adagrad|adadelta|adam]")
    parser.add_argument("-lr", type=float, default=1e-3, help="Initial learning rate")
    parser.add_argument('-max_grad_norm', type=float, default=5, help="""If the norm of the gradient vector exceeds this,
                        renormalize it to have the norm equal to max_grad_norm""")
    parser.add_argument('-dropout', type=float, default=0.3, help='Dropout probability; applied between LSTM stacks.')

    parser.add_argument('-learning_rate_decay', type=float, default=0.5,
                        help="""If update_learning_rate, decay learning rate by
                        this much if (i) perplexity does not decrease on the
                        validation set and (ii) epoch has gone past start_decay_at""")
    parser.add_argument('-start_decay_at', type=int, default=5,
                        help="""Start decaying every epoch after and including this epoch""")

    # GPU
    parser.add_argument('-gpus', default=[0], nargs='+', type=int, help="Use CUDA on the listed devices.")
    parser.add_argument('-log_interval', type=int, default=50, help="Print stats at this interval.")
    parser.add_argument('-seed', type=int, default=3435, # default=-1
                        help="""Random seed used for the experiments reproducibility.""")
    # Critic
    parser.add_argument("-start_reinforce", type=int, default=None, help="""Epoch to start reinforcement training. Use -1 to start immediately.""")
    parser.add_argument("-critic_pretrain_epochs", type=int, default=0, help="Number of epochs to pretrain critic (actor fixed).")
    parser.add_argument("-reinforce_lr", type=float, default=1e-4, help="""Learning rate for reinforcement training.""")

    # Generation
    parser.add_argument("-max_predict_length", required=True, type=int, default=20, help="Maximum length of predictions.")
    parser.add_argument("-predict_mask", type=int, default=0, help="Set to 1 for avoiding repeatitive words and UNK in eval.")

    # Evaluation
    parser.add_argument("-eval", action="store_true", help="Evaluate model only")
    parser.add_argument("-sent_reward", default="bleu", choices=["cr", "bleu"], help="Sentence reward.")
    parser.add_argument("-eval_codenn", action="store_true", help="Set to True to evaluate on codenn DEV/EVAL. Used for evaluation only.")
    parser.add_argument("-eval_codenn_all", action="store_true",
                        help="Set to True to evaluate on codenn test set. Used for evaluation only.")
    parser.add_argument("-collect_anno", action="store_true", help="Set to True to collect generated annotations.")

    opt = parser.parse_args()
    opt.iteration = 0
    return opt

def load_data(opt):
    dataset = torch.load(opt.data)
    dicts = dataset["dicts"]

    supervised_data = lib.Dataset(dataset["train"], "sl_train", opt.batch_size, opt.cuda, eval=False)
    rl_data = lib.Dataset(dataset["train"], "rl_train", opt.batch_size, opt.cuda, eval=False)
    valid_data = lib.Dataset(dataset["valid"], "val", 50, opt.cuda, eval=True) #opt.batch_size
    test_data = lib.Dataset(dataset["test"], "test", 50, opt.cuda, eval=True)

    if "DEV" in dataset:
        DEV = lib.Dataset(dataset['DEV'], "DEV", opt.batch_size, opt.cuda, eval=True)
        EVAL = lib.Dataset(dataset['EVAL'], "EVAL", opt.batch_size, opt.cuda, eval=True)
    else:
        DEV = None
        EVAL = None

    print(" * vocabulary size. source = %d; target = %d" % (dicts["src"].size(), dicts["tgt"].size()))
    print(" * number of XENT training sentences. %d" % len(dataset["train"]["src"]))
    print(" * number of PG training sentences. %d" % len(dataset["train"]["src"]))
    print(" * number of val sentences. %d" % len(dataset["valid"]["src"]))
    print(" * number of test sentences. %d" % len(dataset["test"]["src"]))
    if "DEV" in dataset:
        print(" * number of DEV sentences. %d" % len(dataset["DEV"]["src"]))
        print(" * number of EVAL sentences. %d" % len(dataset["EVAL"]["src"]))
    print(" * maximum batch size. %d" % opt.batch_size)

    return dicts, supervised_data, rl_data, valid_data, test_data, DEV, EVAL

def get_aligned_embedding(emb_old, dict):
    """
    Get an aligned embedding. Missing values will be randomly initialized.
    :param emb_old: a matrix of shape [vocab_size, vec_dim].
    :param dict: a Dict type of dictionary.
    :return:
    """
    w2v = emb_old.wv
    print("The pretrained emb matrix contains %d words, while the given dict contains %d words..." % (
        len(w2v.vocab), dict.size()))

    emb = []
    for idx, word in dict.idxToLabel.items():
        if word in w2v:
            emb.append(w2v[word])
        else:
            emb.append(np.random.uniform(-opt.param_init, opt.param_init, opt.word_vec_size))

    emb = torch.Tensor(emb)
    if opt.cuda:
        emb = emb.cuda()

    return emb


def init(model, dicts):
    for p in model.parameters():
        p.data.uniform_(-opt.param_init, opt.param_init)

    if opt.load_embedding_from is not None:
        emb_src = gensim.models.Word2Vec.load(opt.load_embedding_from + '%s.processed_all.train_xe.src.gz' % opt.lang)
        aligned_emb_src = get_aligned_embedding(emb_src, dicts['src'])
        emb_tgt = gensim.models.Word2Vec.load(opt.load_embedding_from + '%s.processed_all.train_xe.tgt.gz' % opt.lang)
        aligned_emb_tgt = get_aligned_embedding(emb_tgt, dicts['tgt'])
        print("Loading pretrained W2V...")
        pretrained_params = {"encoder.word_lut.weight": aligned_emb_src,
                             "decoder.word_lut.weight": aligned_emb_tgt}
        model.load_state_dict(pretrained_params, strict=False)


def create_optim(model):
    optim = lib.Optim(
        model.parameters(), opt.optim, opt.lr, opt.max_grad_norm,
        lr_decay=opt.learning_rate_decay, start_decay_at=opt.start_decay_at
    )
    return optim

def create_model(model_class, dicts, gen_out_size):
    encoder = lib.Encoder(opt, dicts["src"])
    decoder = lib.TreeDecoder(opt, dicts["tgt"])
    # Use memory efficient generator when output size is large and
    # max_generator_batches is smaller than batch_size.
    if opt.max_generator_batches < opt.batch_size and gen_out_size > 1:
        generator = lib.MemEfficientGenerator(nn.Linear(opt.rnn_size, gen_out_size), opt)
    else:
        generator = lib.BaseGenerator(nn.Linear(opt.rnn_size, gen_out_size), opt)
    model = model_class(encoder, decoder, generator, opt)
    init(model, dicts)
    optim = create_optim(model)

    return model, optim

def create_critic(checkpoint, dicts, opt):
    if opt.load_from is not None and "critic" in checkpoint:
        critic = checkpoint["critic"]
        critic_optim = checkpoint["critic_optim"]
    else:
        critic, critic_optim = create_model(lib.Seq2SeqModel, dicts, 1)
    if opt.cuda:
        critic.cuda(opt.gpus[0])
    return critic, critic_optim


def main():
    print("Start...")
    global opt
    opt = get_opt()

    # Set seed
    torch.manual_seed(opt.seed)
    np.random.seed(opt.seed)
    random.seed(opt.seed)

    opt.cuda = torch.cuda.is_available() and len(opt.gpus)

    if opt.save_dir and not os.path.exists(opt.save_dir):
        os.makedirs(opt.save_dir)

    if torch.cuda.is_available() and not opt.cuda:
        print("WARNING: You have a CUDA device, so you should probably run with -gpus 1")

    if opt.cuda:
        cuda.set_device(opt.gpus[0])
        torch.cuda.manual_seed(opt.seed)

    dicts, supervised_data, rl_data, valid_data, test_data, DEV, EVAL = load_data(opt)

    print("Building model...")

    use_critic = opt.start_reinforce is not None
    print("use_critic: ", use_critic)
    print("has_baseline: ", opt.has_baseline)
    if not opt.has_baseline:
        assert opt.critic_pretrain_epochs == 0

    if opt.load_from is None:
        model, optim = create_model(lib.Seq2SeqModel, dicts, dicts["tgt"].size())
        checkpoint = None

    else:
        print("Loading from checkpoint at %s" % opt.load_from)
        checkpoint = torch.load(opt.load_from)#, map_location=lambda storage, loc: storage)
        model = checkpoint["model"]
        # config testing
        for attribute in ["predict_mask", "max_predict_length"]:
            model.opt.__dict__[attribute] = opt.__dict__[attribute]
        optim = checkpoint["optim"]
        optim.start_decay_at = opt.start_decay_at
        if optim.start_decay_at > opt.end_epoch:
            print("No decay!")
        opt.start_epoch = checkpoint["epoch"] + 1

    print("model: ", model)
    print("optim: ", optim)

    # GPU.
    if opt.cuda:
        model.cuda(opt.gpus[0])

    # Start reinforce training immediately.
    print("opt.start_reinforce: ", opt.start_reinforce)

    # Check if end_epoch is large enough.
    if use_critic:
        assert opt.start_epoch + opt.critic_pretrain_epochs - 1 <= \
               opt.end_epoch, "Please increase -end_epoch to perform pretraining!"

    nParams = sum([p.nelement() for p in model.parameters()])
    print("* number of parameters: %d" % nParams)

    if opt.sent_reward == "cr":
        lib.RetReward.cr = code_retrieval.CrCritic()

    # Metrics.
    print("sent_reward: %s" % opt.sent_reward)
    metrics = {}
    metrics["xent_loss"] = lib.Loss.weighted_xent_loss
    metrics["critic_loss"] = lib.Loss.weighted_mse
    if opt.sent_reward == "bleu":
        metrics["sent_reward"] = {"train": lib.Reward.wrapped_sentence_bleu,
                                  "eval": lib.Reward.wrapped_sentence_bleu}
    else:
        metrics["sent_reward"] = {"train": lib.RetReward.retrieval_mrr_train,
                                  "eval": lib.RetReward.retrieval_mrr_eval}

    print("opt.eval: ", opt.eval)
    print("opt.eval_codenn: ", opt.eval_codenn)
    print("opt.eval_codenn_all: ", opt.eval_codenn_all)
    print("opt.collect_anno: ", opt.collect_anno)

    # Evaluate model
    if opt.eval:
        if opt.sent_reward == "cr" and (opt.eval_codenn or opt.eval_codenn_all):
            raise Exception("Currently we do not support evaluating MRR on codenn!")

        if False:
            # On training set.
            if opt.sent_reward == "cr":
                metrics["sent_reward"]["eval"] = lib.RetReward.retrieval_mrr_train
            #if opt.collect_anno:
            #    metrics["sent_reward"] = {"train": None, "eval": None}

            evaluator = lib.Evaluator(model, metrics, dicts, opt)
            pred_file = opt.load_from.replace(".pt", ".train.pred")
            if opt.eval_codenn or opt.eval_codenn_all:
                raise Exception("Invalid eval_codenn!")
            print("train_data.src: ", len(supervised_data.src))
            if opt.predict_mask:
                pred_file += ".masked"
            pred_file += ".metric%s" % opt.sent_reward
            evaluator.eval(supervised_data, pred_file)

        if True:
            # On validation set.
            if opt.sent_reward == "cr":
                metrics["sent_reward"]["eval"] = lib.RetReward.retrieval_mrr_eval
            #if opt.collect_anno:
            #    metrics["sent_reward"] = {"train": None, "eval": None}

            evaluator = lib.Evaluator(model, metrics, dicts, opt)
            pred_file = opt.load_from.replace(".pt", ".valid.pred")
            if opt.eval_codenn:
                pred_file = pred_file.replace("valid", "DEV")
                valid_data = DEV
            elif opt.eval_codenn_all:
                pred_file = pred_file.replace("valid", "DEV_all")
                print("* Please input valid data = DEV_all")
            print("valid_data.src: ", len(valid_data.src))
            if opt.predict_mask:
                pred_file += ".masked"
            pred_file += ".metric%s" % opt.sent_reward
            evaluator.eval(valid_data, pred_file)

        if False:
            # On test set.
            if opt.sent_reward == "cr":
                metrics["sent_reward"]["eval"] = lib.RetReward.retrieval_mrr_eval
            #if opt.collect_anno:
            #    metrics["sent_reward"] = {"train": None, "eval": None}

            evaluator = lib.Evaluator(model, metrics, dicts, opt)
            pred_file = opt.load_from.replace(".pt", ".test.pred")
            if opt.eval_codenn:
                pred_file = pred_file.replace("test", "EVAL")
                test_data = EVAL
            elif opt.eval_codenn_all:
                pred_file = pred_file.replace("test", "EVAL_all")
                print("* Please input test data = EVAL_all")
            print("test_data.src: ", len(test_data.src))
            if opt.predict_mask:
                pred_file += ".masked"
            pred_file += ".metric%s" % opt.sent_reward
            evaluator.eval(test_data, pred_file)

    else:
        print("supervised_data.src: ", len(supervised_data.src))
        print("supervised_data.tgt: ", len(supervised_data.tgt))
        xent_trainer = lib.Trainer(model, supervised_data, valid_data, metrics, dicts, optim, opt, DEV=DEV)

        if use_critic:
            start_time = time.time()
            # Supervised training.
            print("supervised training..")
            print("start_epoch: ", opt.start_epoch)

            xent_trainer.train(opt.start_epoch, opt.start_reinforce - 1, start_time)

            if opt.sent_reward == "bleu":
                _valid_data = DEV
            else:
                _valid_data = valid_data

            if opt.has_baseline:
                # Create critic here to not affect random seed.
                critic, critic_optim = create_critic(checkpoint, dicts, opt)
                print("Building critic...")
                print("Critic: ", critic)
                print("Critic optim: ", critic_optim)

                # Pretrain critic.
                print("pretrain critic...")
                if opt.critic_pretrain_epochs > 0:
                    reinforce_trainer = lib.ReinforceTrainer(model, critic, supervised_data, _valid_data, metrics, dicts, optim, critic_optim, opt)
                    reinforce_trainer.train(opt.start_reinforce, opt.start_reinforce + opt.critic_pretrain_epochs - 1, True, start_time)
            else:
                print("NOTE: do not have a baseline model")
                critic, critic_optim = None, None

            # Reinforce training.
            print("reinforce training...")
            reinforce_trainer = lib.ReinforceTrainer(model, critic, rl_data, _valid_data, metrics, dicts, optim, critic_optim, opt)
            reinforce_trainer.train(opt.start_reinforce + opt.critic_pretrain_epochs, opt.end_epoch, False, start_time)

        else: # Supervised training only. Set opt.start_reinforce to None
            xent_trainer.train(opt.start_epoch, opt.end_epoch)


if __name__ == '__main__':
    main()