train_rmc.py

# copypasta from main.py of pytorch word_language_model code
# coding: utf-8
import argparse
import time
import math
import os
import torch
import torch.nn as nn
import torch.onnx
import datetime
import shutil
import pickle
import data
from relational_rnn_models import RelationalMemory

# is it faster?
torch.backends.cudnn.benchmark = True

parser = argparse.ArgumentParser(description='PyTorch Wikitext-2 RNN/LSTM Language Model')
# hyperparams for text data
parser.add_argument('--data', type=str, default='./data/wikitext-2',
                    help='location of the data corpus')
parser.add_argument('--emsize', type=int, default=192,
                    help='size of word embeddings')

# NEW!: hyperparams for relational memory core (RMC)
parser.add_argument('--memslots', type=int, default=1,
                    help='number of memory slots of the relational memory core')
parser.add_argument('--headsize', type=int, default=192,
                    help='size of the each head for multihead attention')
parser.add_argument('--numheads', type=int, default=4,
                    help='total number of heads for multihead attention')
parser.add_argument('--numblocks', type=int, default=1,
                    help='Number of times to compute attention per time step')
parser.add_argument('--forgetbias', type=float, default=1.,
                    help='Bias to use for the forget gate, assuming we are using some form of gating')
parser.add_argument('--inputbias', type=float, default=0.,
                    help='Bias to use for the input gate, assuming we are using some form of gating')
parser.add_argument('--gatestyle', type=str, default='unit',
                    help='Whether to use per-element gating (\'unit\'), per-memory slot gating (\'memory\'), or no gating at all (None).')
parser.add_argument('--attmlplayers', type=int, default=3,
                    help='Number of layers to use in the post-attention MLP')
parser.add_argument('--keysize', type=int, default=64,
                    help='Size of vector to use for key & query vectors in the attention'
                         'computation. Defaults to None, in which case we use `head_size`')
# parameters for adaptive softmax
parser.add_argument('--adaptivesoftmax', action='store_true',
                    help='use adaptive softmax during hidden state to output logits.'
                         'it uses less memory by approximating softmax of large vocabulary.')
parser.add_argument('--cutoffs', nargs="*", type=int, default=[10000, 50000, 100000],
                    help='cutoff values for adaptive softmax. list of integers.'
                         'optimal values are based on word frequencey and vocabulary size of the dataset.')

# other hyperparams for general RNN mechanics
parser.add_argument('--lr', type=float, default=0.001,
                    help='initial learning rate')
parser.add_argument('--clip', type=float, default=0.1,
                    help='gradient clipping')
parser.add_argument('--epochs', type=int, default=100,
                    help='upper epoch limit')
parser.add_argument('--batch_size', type=int, default=64, metavar='N',
                    help='batch size')
parser.add_argument('--bptt', type=int, default=100,
                    help='sequence length')
# dropout of RMC is hard-bound to 0.5 at the embedding layer
# parser.add_argument('--dropout', type=float, default=0.2,
#                     help='dropout applied to layers (0 = no dropout)')
# embed weight tying is set always to true
# parser.add_argument('--tied', action='store_true',
#                     help='tie the word embedding and softmax weights')
parser.add_argument('--seed', type=int, default=1111,
                    help='random seed')
parser.add_argument('--cuda', action='store_true',
                    help='use CUDA')
parser.add_argument('--log-interval', type=int, default=100, metavar='N',
                    help='report interval')
parser.add_argument('--onnx-export', type=str, default='',
                    help='path to export the final model in onnx format')
parser.add_argument('--resume', type=int, default=None,
                    help='if specified with the 1-indexed global epoch, loads the checkpoint and resumes training')

# experiment name for this run
parser.add_argument('--name', type=str, default=None,
                    help='name for this experiment. generates folder with the name if specified.')

args = parser.parse_args()

# Set the random seed manually for reproducibility.
torch.manual_seed(args.seed)

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

device = torch.device("cuda" if args.cuda else "cpu")
###############################################################################
# Load data
###############################################################################
corpus_name = os.path.basename(os.path.normpath(args.data))
corpus_filename = './data/corpus-' + str(corpus_name) + str('.pkl')
if os.path.isfile(corpus_filename):
    print("loading pre-built " + str(corpus_name) + " corpus file...")
    loadfile = open(corpus_filename, 'rb')
    corpus = pickle.load(loadfile)
    loadfile.close()
else:
    print("building " + str(corpus_name) + " corpus...")
    corpus = data.Corpus(args.data)
    # save the corpus for later
    savefile = open(corpus_filename, 'wb')
    pickle.dump(corpus, savefile)
    savefile.close()
    print("corpus saved to pickle")


# Starting from sequential data, batchify arranges the dataset into columns.
# For instance, with the alphabet as the sequence and batch size 4, we'd get
# ┌ a g m s ┐
# │ b h n t │
# │ c i o u │
# │ d j p v │
# │ e k q w │
# └ f l r x ┘.
# These columns are treated as independent by the model, which means that the
# dependence of e. g. 'g' on 'f' can not be learned, but allows more efficient
# batch processing.

def batchify(data, bsz):
    # Work out how cleanly we can divide the dataset into bsz parts.
    nbatch = data.size(0) // bsz
    # Trim off any extra elements that wouldn't cleanly fit (remainders).
    data = data.narrow(0, 0, nbatch * bsz)
    # Evenly divide the data across the bsz batches.
    data = data.view(bsz, -1).t().contiguous()
    return data.to(device)


eval_batch_size = 32
train_data = batchify(corpus.train, args.batch_size)
val_data = batchify(corpus.valid, eval_batch_size)
test_data = batchify(corpus.test, eval_batch_size)

# create folder for current experiments
# name: args.name + current time
# includes: entire scripts for faithful reproduction, train & test logs
folder_name = str(datetime.datetime.now())[:-7]
if args.name is not None:
    folder_name = str(args.name) + ' ' + folder_name

os.mkdir(folder_name)
for file in os.listdir(os.getcwd()):
    if file.endswith(".py"):
        shutil.copy2(file, os.path.join(os.getcwd(), folder_name))
logger_train = open(os.path.join(os.getcwd(), folder_name, 'train_log.txt'), 'w+')
logger_test = open(os.path.join(os.getcwd(), folder_name, 'test_log.txt'), 'w+')

# save args to logger
logger_train.write(str(args) + '\n')

# define saved model file location
savepath = os.path.join(os.getcwd(), folder_name)

###############################################################################
# Build the model
###############################################################################

ntokens = len(corpus.dictionary)
print("vocabulary size (ntokens): " + str(ntokens))
if args.adaptivesoftmax:
    print("Adaptive Softmax is on: the performance depends on cutoff values. check if the cutoff is properly set")
    print("Cutoffs: " + str(args.cutoffs))
    if args.cutoffs[-1] > ntokens:
        raise ValueError("the last element of cutoff list must be lower than vocab size of the dataset")

model = RelationalMemory(mem_slots=args.memslots, head_size=args.headsize, input_size=args.emsize, num_tokens=ntokens,
                         num_heads=args.numheads, num_blocks=args.numblocks, forget_bias=args.forgetbias,
                         input_bias=args.inputbias, attention_mlp_layers=args.attmlplayers, key_size=args.keysize,
                         use_adaptive_softmax=args.adaptivesoftmax, cutoffs=args.cutoffs).to(device)
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
model = nn.DataParallel(model)

optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=5)

###############################################################################
# Load the model checkpoint if specified and restore the global & best epoch
###############################################################################
if args.resume is not None:
    print("--resume detected. loading checkpoint...")
global_epoch = args.resume if args.resume is not None else 0
best_epoch = args.resume if args.resume is not None else 0
if args.resume is not None:
    loadpath = os.path.join(os.getcwd(), "model_{}.pt".format(args.resume))
    if not os.path.isfile(loadpath):
        raise FileNotFoundError(
            "model_{}.pt not found. place the model checkpoint file to the current working directory.".format(
                args.resume))
    checkpoint = torch.load(loadpath)
    model.load_state_dict(checkpoint["state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    global_epoch = checkpoint["global_epoch"]
    best_epoch = checkpoint["best_epoch"]

print("model built, total trainable params: " + str(total_params))


###############################################################################
# Training code
###############################################################################

# get_batch subdivides the source data into chunks of length args.bptt.
# If source is equal to the example output of the batchify function, with
# a bptt-limit of 2, we'd get the following two Variables for i = 0:
# ┌ a g m s ┐ ┌ b h n t ┐
# └ b h n t ┘ └ c i o u ┘
# Note that despite the name of the function, the subdivison of data is not
# done along the batch dimension (i.e. dimension 1), since that was handled
# by the batchify function. The chunks are along dimension 0, corresponding
# to the seq_len dimension in the LSTM.


def get_batch(source, i):
    seq_len = min(args.bptt, len(source) - 1 - i)
    data = source[i:i + seq_len]
    target = source[i + 1:i + 1 + seq_len].view(-1)
    return data, target


def evaluate(data_source):
    # Turn on evaluation mode which disables dropout.
    model.eval()
    total_loss = 0.
    ntokens = len(corpus.dictionary)
    memory = model.module.initial_state(eval_batch_size, trainable=False).to(device)

    with torch.no_grad():
        for i in range(0, data_source.size(0) - 1, args.bptt):
            data, targets = get_batch(data_source, i)
            data = torch.t(data)

            loss, memory = model(data, memory, targets)
            loss = torch.mean(loss)

            # data has shape [T * B, N]
            total_loss += args.bptt * loss.item()

    return total_loss / len(data_source)


def train():
    # Turn on training mode which enables dropout.
    model.train()
    total_loss = 0.
    forward_elapsed_time = 0.
    start_time = time.time()
    ntokens = len(corpus.dictionary)
    # in RMC, "hidden state" is called "memory" instead. so use the name "memory"
    memory = model.module.initial_state(args.batch_size, trainable=True).to(device)

    for batch, i in enumerate(range(0, train_data.size(0) - 1, args.bptt)):
        data, targets = get_batch(train_data, i)
        # transpose the data to [batch, seq]
        data = torch.t(data)

        # synchronize cuda for a proper speed benchmark
        torch.cuda.synchronize()

        forward_start_time = time.time()
        model.zero_grad()

        # the forward pass of RMC just returns loss and does not return logits (DataParallel code optimization)
        loss, memory = model(data, memory, targets)
        loss = torch.mean(loss)
        total_loss += loss.item()

        # synchronize cuda for a proper speed benchmark
        torch.cuda.synchronize()

        forward_elapsed = time.time() - forward_start_time
        forward_elapsed_time += forward_elapsed

        loss.backward()

        # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)

        optimizer.step()

        if batch % args.log_interval == 0 and batch > 0:
            cur_loss = total_loss / args.log_interval
            elapsed = time.time() - start_time
            printlog = '| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.5f} | ms/batch {:5.2f} | forward ms/batch {:5.2f} | loss {:5.2f} | ppl {:8.2f}'.format(
                epoch, batch, len(train_data) // args.bptt, optimizer.param_groups[0]['lr'],
                              elapsed * 1000 / args.log_interval, forward_elapsed_time * 1000 / args.log_interval,
                cur_loss, math.exp(cur_loss))
            # print and save the log
            print(printlog)
            logger_train.write(printlog + '\n')
            logger_train.flush()
            total_loss = 0.
            # reset timer
            start_time = time.time()
            forward_start_time = time.time()
            forward_elapsed_time = 0.


def export_onnx(path, batch_size, seq_len):
    print('The model is also exported in ONNX format at {}'.
          format(os.path.realpath(args.onnx_export)))
    model.eval()
    dummy_input = torch.LongTensor(seq_len * batch_size).zero_().view(-1, batch_size).to(device)
    hidden = model.init_hidden(batch_size)
    torch.onnx.export(model, (dummy_input, hidden), path)


# Loop over epochs.
best_val_loss = None

# At any point you can hit Ctrl + C to break out of training early.
try:
    print("training started...")
    if global_epoch > args.epochs:
        raise ValueError("global_epoch is higher than args.epochs when resuming training.")
    for epoch in range(global_epoch + 1, args.epochs + 1):
        global_epoch += 1
        epoch_start_time = time.time()
        train()
        val_loss = evaluate(val_data)

        print('-' * 89)
        testlog = '| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | valid ppl {:8.2f}'.format(epoch, (
                time.time() - epoch_start_time), val_loss, math.exp(val_loss))
        print(testlog)
        logger_test.write(testlog + '\n')
        logger_test.flush()
        print('-' * 89)

        scheduler.step(val_loss)

        # Save the model if the validation loss is the best we've seen so far.
        # model_{} contains state_dict and other states, model_dump_{} contains all the dependencies for generate_rmc.py
        if not best_val_loss or val_loss < best_val_loss:
            try:
                os.remove(os.path.join(savepath, "model_{}.pt".format(best_epoch)))
                os.remove(os.path.join(savepath, "model_dump_{}.pt").format(best_epoch))
            except FileNotFoundError:
                pass
            best_epoch = global_epoch
            torch.save(model, os.path.join(savepath, "model_dump_{}.pt".format(global_epoch)))
            with open(os.path.join(savepath, "model_{}.pt".format(global_epoch)), 'wb') as f:
                optimizer_state = optimizer.state_dict()
                scheduler_state = scheduler.state_dict()
                torch.save({"state_dict": model.state_dict(),
                            "optimizer": optimizer_state,
                            "scheduler": scheduler_state,
                            "global_epoch": global_epoch,
                            "best_epoch": best_epoch}, f)
            best_val_loss = val_loss
        else:
            pass

except KeyboardInterrupt:
    print('-' * 89)
    print('Exiting from training early: loading checkpoint from the best epoch {}...'.format(best_epoch))

# Load the best saved model.
with open(os.path.join(savepath, "model_{}.pt".format(best_epoch)), 'rb') as f:
    checkpoint = torch.load(f)
    model.load_state_dict(checkpoint["state_dict"])

# Run on test data.
test_loss = evaluate(test_data)

print('=' * 89)
testlog = '| End of training | test loss {:5.2f} | test ppl {:8.2f}'.format(
    test_loss, math.exp(test_loss))
print(testlog)
logger_test.write(testlog + '\n')
logger_test.flush()
print('=' * 89)

if len(args.onnx_export) > 0:
    # Export the model in ONNX format.
    export_onnx(args.onnx_export, batch_size=1, seq_len=args.bptt)