train_mulitple_datasets.py

import os
import json
import time
import torch
import argparse
import logging
import numpy as np
from multiprocessing import cpu_count
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader
from collections import OrderedDict, defaultdict

from ptb import PTB, Brown
from utils import to_var, idx2word, experiment_name
from model import SentenceVAE

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)


def main(args):
    ts = time.strftime('%Y-%b-%d-%H:%M:%S', time.gmtime())

    splits = ['train', 'valid'] + (['test'] if args.test else [])

    datasets = OrderedDict()
    for split in splits:
        datasets[split] = [
            PTB(
                data_dir=args.data_dir,
                split=split,
                create_data=args.create_data,
                max_sequence_length=args.max_sequence_length,
                min_occ=args.min_occ
            ),
            Brown(
                data_dir=args.data_dir,
                split=split,
                create_data=args.create_data,
                max_sequence_length=args.max_sequence_length,
                min_occ=args.min_occ
            )]

    for split in splits:
        ptb, brown = datasets[split]

        # Make sure the two datasets are of the same length
        dataset_length = min(len(ptb), len(brown))
        if len(ptb) != len(brown):
            ptb.newlength(dataset_length)
            brown.newlength(dataset_length)

    model = SentenceVAE(
        vocab_size=datasets['train'][0].vocab_size,
        sos_idx=datasets['train'][0].sos_idx,
        eos_idx=datasets['train'][0].eos_idx,
        pad_idx=datasets['train'][0].pad_idx,
        unk_idx=datasets['train'][0].unk_idx,
        max_sequence_length=args.max_sequence_length,
        embedding_size=args.embedding_size,
        rnn_type=args.rnn_type,
        hidden_size=args.hidden_size,
        word_dropout=args.word_dropout,
        embedding_dropout=args.embedding_dropout,
        latent_size=args.latent_size,
        num_layers=args.num_layers,
        bidirectional=args.bidirectional
    )

    if torch.cuda.is_available():
        model = model.cuda()

    print(model)

    if args.tensorboard_logging:
        writer = SummaryWriter(os.path.join(args.logdir, experiment_name(args, ts)))
        writer.add_text("model", str(model))
        writer.add_text("args", str(args))
        writer.add_text("ts", ts)

    save_model_path = os.path.join(args.save_model_path, ts)
    os.makedirs(save_model_path)

    total_steps = ((len(datasets["train"][0]) + len(datasets["train"][1])) // args.batch_size) * args.epochs
    print("Train dataset size", total_steps)

    def kl_anneal_function(anneal_function, step):
        if anneal_function == 'identity':
            return 1
        if anneal_function == 'linear':
            if args.warmup is None:
                return 1 - (total_steps - step) / total_steps
            else:
                warmup_steps = (total_steps / args.epochs) * args.warmup
                return 1 - (warmup_steps - step) / warmup_steps if step < warmup_steps else 1.0

    ReconLoss = torch.nn.NLLLoss(size_average=False, ignore_index=datasets['train'][0].pad_idx)

    def loss_fn(logp, target, length, mean, logv, anneal_function, step):

        # cut-off unnecessary padding from target, and flatten
        target = target[:, :torch.max(length).data[0]].contiguous().view(-1)
        logp = logp.view(-1, logp.size(2))

        # Negative Log Likelihood
        recon_loss = ReconLoss(logp, target)

        # KL Divergence
        KL_loss = -0.5 * torch.sum(1 + logv - mean.pow(2) - logv.exp())
        KL_weight = kl_anneal_function(anneal_function, step)

        return recon_loss, KL_loss, KL_weight

    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)

    tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.Tensor
    step = 0
    for epoch in range(args.epochs):

        for split in splits:

            data_loaders = [DataLoader(
                dataset=datasets[split][0],
                batch_size=args.batch_size,
                shuffle=split == 'train',
                num_workers=cpu_count(),
                pin_memory=torch.cuda.is_available()
            ), DataLoader(
                dataset=datasets[split][1],
                batch_size=args.batch_size,
                shuffle=split == 'train',
                num_workers=cpu_count(),
                pin_memory=torch.cuda.is_available()
            )]

            tracker = defaultdict(tensor)

            # Enable/Disable Dropout
            if split == 'train':
                model.train()
            else:
                model.eval()

            total_length = len(data_loaders[0]) + len(data_loaders[1])
            iteration = -1
            for data_loader in data_loaders:
                for batch in data_loader:
                    iteration += 1
                    batch_size = batch['input'].size(0)

                    for k, v in batch.items():
                        if torch.is_tensor(v):
                            batch[k] = to_var(v)

                    # Forward pass
                    logp, mean, logv, z = model(batch['input'], batch['length'])

                    # loss calculation
                    recon_loss, KL_loss, KL_weight = loss_fn(logp, batch['target'],
                                                             batch['length'], mean, logv, args.anneal_function, step)

                    if split == 'train':
                        loss = (recon_loss + KL_weight * KL_loss) / batch_size
                    else:
                        # report complete elbo when validation
                        loss = (recon_loss + KL_loss) / batch_size

                    # backward + optimization
                    if split == 'train':
                        optimizer.zero_grad()
                        loss.backward()
                        optimizer.step()
                        step += 1

                    # bookkeepeing
                    tracker['negELBO'] = torch.cat((tracker['negELBO'], loss.data.unsqueeze(0)))

                    if args.tensorboard_logging:
                        writer.add_scalar("%s/Negative_ELBO" % split.upper(), loss.data[0],
                                          epoch * total_length + iteration)
                        writer.add_scalar("%s/Recon_Loss" % split.upper(), recon_loss.data[0] / batch_size,
                                          epoch * total_length + iteration)
                        writer.add_scalar("%s/KL_Loss" % split.upper(), KL_loss.data[0] / batch_size,
                                          epoch * total_length + iteration)
                        writer.add_scalar("%s/KL_Weight" % split.upper(), KL_weight,
                                          epoch * total_length + iteration)

                    if iteration % args.print_every == 0 or iteration + 1 == len(data_loader):
                        logger.info("%s Batch %04d/%i, Loss %9.4f, Recon-Loss %9.4f, KL-Loss %9.4f, KL-Weight %6.3f"
                                    % (split.upper(), iteration, total_length - 1, loss.data[0],
                                       recon_loss.data[0] / batch_size, KL_loss.data[0] / batch_size, KL_weight))

                    if split == 'valid':
                        if 'target_sents' not in tracker:
                            tracker['target_sents'] = list()
                        tracker['target_sents'] += idx2word(batch['target'].data, i2w=datasets['train'][0].get_i2w(),
                                                            pad_idx=datasets['train'][0].pad_idx)

                        tracker['z'] = torch.cat((tracker['z'], z.data), dim=0)
                        tracker['dataset'] = torch.cat((tracker['dataset'], torch.tensor(batch['dataset']).float()), dim=0)

            logger.info("%s Epoch %02d/%i, Mean Negative ELBO %9.4f" % (
                split.upper(), epoch, args.epochs, torch.mean(tracker['negELBO'])))

            if args.tensorboard_logging:
                writer.add_scalar("%s-Epoch/NegELBO" % split.upper(), torch.mean(tracker['negELBO']), epoch)

            # save a dump of all sentences and the encoded latent space
            if split == 'valid':
                dump = {'target_sents': tracker['target_sents'], 'z': tracker['z'].tolist(), "dataset": tracker["dataset"].tolist()}
                if not os.path.exists(os.path.join('dumps', ts)):
                    os.makedirs('dumps/' + ts)
                with open(os.path.join('dumps/' + ts + '/valid_E%i.json' % epoch), 'w') as dump_file:
                    json.dump(dump, dump_file)

            # save checkpoint
            if split == 'train':
                checkpoint_path = os.path.join(save_model_path, "E%i.pytorch" % (epoch))
                torch.save(model.state_dict(), checkpoint_path)
                logger.info("Model saved at %s" % checkpoint_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--data_dir', type=str, default='data')
    parser.add_argument('--create_data', action='store_true')
    parser.add_argument('--max_sequence_length', type=int, default=60)
    parser.add_argument('--min_occ', type=int, default=1)
    parser.add_argument('--test', action='store_true')

    parser.add_argument('-ep', '--epochs', type=int, default=10)
    parser.add_argument('-bs', '--batch_size', type=int, default=32)
    parser.add_argument('-lr', '--learning_rate', type=float, default=0.001)

    parser.add_argument('-eb', '--embedding_size', type=int, default=300)
    parser.add_argument('-rnn', '--rnn_type', type=str, default='gru')
    parser.add_argument('-hs', '--hidden_size', type=int, default=256)
    parser.add_argument('-nl', '--num_layers', type=int, default=1)
    parser.add_argument('-bi', '--bidirectional', action='store_true')
    parser.add_argument('-ls', '--latent_size', type=int, default=16)
    parser.add_argument('-wd', '--word_dropout', type=float, default=0)
    parser.add_argument('-ed', '--embedding_dropout', type=float, default=0.5)

    parser.add_argument('-af', '--anneal_function', type=str, default='identity')
    parser.add_argument('-wu', '--warmup', type=int, default=None)

    parser.add_argument('-v', '--print_every', type=int, default=50)
    parser.add_argument('-tb', '--tensorboard_logging', action='store_true')
    parser.add_argument('-log', '--logdir', type=str, default='logs')
    parser.add_argument('-bin', '--save_model_path', type=str, default='bin')

    args = parser.parse_args()

    args.rnn_type = args.rnn_type.lower()
    args.anneal_function = args.anneal_function.lower()

    assert args.rnn_type in ['rnn', 'lstm', 'gru']
    assert 0 <= args.word_dropout <= 1

    main(args)