train.py

import os
import torch
import random
import traceback
import model.utils as utils
import model.dataset as dataset
from model.model_multi_input import MultiInputModel
from model.trainer_multi_input import Trainer
from model.text import Vocab
import re
from torch.nn.parallel import DistributedDataParallel
import argparse


parser = argparse.ArgumentParser()
parser.add_argument('--config', help='config file', default='/home/data/tmp/NLP_Course/GPT_Dialog_model/config.json')
parser.add_argument('--gpu', help='which gpu to use', type=str, default='1')
parser.add_argument("--local_rank", help='used for distributed training', type=int, default=-1)

args = parser.parse_args()
config = utils.load_config(args.config)
config_path = os.path.dirname(args.config)
logger = utils.get_logger(os.path.join(config_path, 'main.log'))

train_dir = os.path.join(config_path, config['train_dir'])
data_dir = os.path.join(config_path, config['data_dir'])
eval_dir = os.path.join(config_path, config['eval_dir'])
log_dir = os.path.join(config_path, config['log_dir'])
best_model = os.path.join(config_path, config['best_dir'])

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu


# helpers -----------------------------------------------------
def save_func(epoch, device):
    filename = utils.get_ckpt_filename('model', epoch)
    torch.save(trainer.state_dict(), os.path.join(train_dir, filename))
    if os.path.exists(os.path.join(train_dir, utils.get_ckpt_filename('model', epoch-80))):
        os.remove(os.path.join(train_dir, utils.get_ckpt_filename('model', epoch-80)))

def sample_text_func(epoch, device):
    n_samples = 8
    samples_idxs = random.sample(range(len(valid_dataset)), n_samples)
    samples = [valid_dataset[idx] for idx in samples_idxs]
    for i, data in enumerate(samples):
        contexts = [torch.tensor([data['post']], dtype=torch.long, device=device)]

        prediction = trainer.model.predict(contexts)[0]
        post_str = vocab.ids2string(data['post'][1:-1])
        resp_str = vocab.ids2string(data['resp'][1:-1])
        pred_str = vocab.ids2string(prediction)

        logger.info('-------epoch {} sample {}---------'.format(epoch, i))
        logger.info('post: {}'.format(post_str))
        logger.info('resp: {}'.format(resp_str))
        logger.info('pred: {}'.format(pred_str))
# helpers -----------------------------------------------------

try:
    if args.local_rank == -1 or args.local_rank == 0:
        logger.info('pytorch version: {}'.format(torch.__version__))
        for i in config:
            logger.info('{}: {}'.format(i, config[i]))
        for i in vars(args):
            logger.info('{}: {}'.format(i, getattr(args, i)))

        dirs = [train_dir, eval_dir, log_dir, best_model]

        for d in dirs:
            if not os.path.isdir(d):
                logger.info('cannot find {}, mkdiring'.format(d))
                os.makedirs(d)

    # code for distributed training
    distributed = (args.local_rank != -1)
    if distributed:
        print(args.local_rank)
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend='nccl', init_method='env://')
        torch.manual_seed(config.seed)
    else:
        device = torch.device("cuda", 0)

    vocab = Vocab(config.vocab_path)
    train_dataset = dataset.DialogDataset([os.path.join(data_dir, config.train_data)],
                                          vocab, logger, config.max_seq_len - 1)
    valid_dataset = dataset.DialogDataset([os.path.join(data_dir, config.valid_data)],
                                          vocab, logger, config.max_seq_len - 1)

    logger.info('Building models')
    model =  MultiInputModel(config, vocab)
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(name, param.shape)

    latest_ckpt = utils.get_latest_ckpt(train_dir)
    if latest_ckpt is None:  # start from scratch
        logger.info('start from CGPT weights')
        cgpt_model = torch.load(config.cgpt_parameters_dir, map_location=device)
        cgpt_model.pop('decoder.pre_softmax.weight')

        b = list(cgpt_model.keys())
        for i in b:
            cgpt_model[i.split('.', 1)[1]] = cgpt_model.pop(i)
        model.transformer_module.load_state_dict(cgpt_model, strict=True)
        logger.info('CGPT weights loaded from {}'.format(config.cgpt_parameters_dir))

    trainer = Trainer(model, train_dataset, valid_dataset, config, log_dir, logger, device, vocab.special_tokens_ids,
                      distributed=distributed)

    if distributed:
        trainer.model.transformer_module = DistributedDataParallel(
            trainer.model.transformer_module, device_ids=[args.local_rank], output_device=args.local_rank)

    start_epoch = 0
    if latest_ckpt is not None:
        logger.info('Weights loading from {}'.format(os.path.join(train_dir, latest_ckpt)))
        start_epoch = utils.get_epoch_from_ckpt(latest_ckpt)
        trainer.load_state_dict(torch.load(os.path.join(train_dir, latest_ckpt), map_location=device))

    try:
        if args.local_rank in [-1, 0]:
            trainer.train(start_epoch, config.n_epochs, after_epoch_funcs=[save_func])
        else:
            trainer.train(start_epoch, config.n_epochs)
        # model_trainer.train(trainer_config.n_epochs, after_epoch_funcs=[sample_text_func], risk_func=f1_risk)
    except (KeyboardInterrupt, Exception, RuntimeError) as e:
        torch.save(trainer.state_dict(), os.path.join(train_dir, 'interrupt.pt'))
        raise e
except:
    logger.error(traceback.format_exc())