train.py

import argparse
import time
import warnings
import zipfile
import os
import collections

os.environ['DGLBACKEND'] = 'mxnet'
os.environ['MXNET_GPU_MEM_POOL_TYPE'] = 'Round'

import numpy as np
import mxnet as mx
from mxnet import gluon

import dgl
import dgl.data as data

from tree_lstm import TreeLSTM

SSTBatch = collections.namedtuple('SSTBatch', ['graph', 'mask', 'wordid', 'label'])

def batcher(ctx):
    def batcher_dev(batch):
        batch_trees = dgl.batch(batch)
        return SSTBatch(graph=batch_trees,
                        mask=batch_trees.ndata['mask'].as_in_context(ctx),
                        wordid=batch_trees.ndata['x'].as_in_context(ctx),
                        label=batch_trees.ndata['y'].as_in_context(ctx))
    return batcher_dev

def prepare_glove():
    if not (os.path.exists('glove.840B.300d.txt')
            and data.utils.check_sha1('glove.840B.300d.txt',
                                      sha1_hash='294b9f37fa64cce31f9ebb409c266fc379527708')):
        zip_path = data.utils.download('http://nlp.stanford.edu/data/glove.840B.300d.zip',
                                       sha1_hash='8084fbacc2dee3b1fd1ca4cc534cbfff3519ed0d')
        with zipfile.ZipFile(zip_path, 'r') as zf:
            zf.extractall()
        if not data.utils.check_sha1('glove.840B.300d.txt',
                                     sha1_hash='294b9f37fa64cce31f9ebb409c266fc379527708'):
            warnings.warn('The downloaded glove embedding file checksum mismatch. File content '
                          'may be corrupted.')

def main(args):
    np.random.seed(args.seed)
    mx.random.seed(args.seed)

    best_epoch = -1
    best_dev_acc = 0

    cuda = args.gpu >= 0
    if cuda:
        if args.gpu in mx.test_utils.list_gpus():
            ctx = mx.gpu(args.gpu)
        else:
            print('Requested GPU id {} was not found. Defaulting to CPU implementation'.format(args.gpu))
            ctx = mx.cpu()
    else:
        ctx = mx.cpu()

    if args.use_glove:
        prepare_glove()

    trainset = data.SSTDataset()
    train_loader = gluon.data.DataLoader(dataset=trainset,
                                         batch_size=args.batch_size,
                                         batchify_fn=batcher(ctx),
                                         shuffle=True,
                                         num_workers=0)
    devset = data.SSTDataset(mode='dev')
    dev_loader = gluon.data.DataLoader(dataset=devset,
                                       batch_size=100,
                                       batchify_fn=batcher(ctx),
                                       shuffle=True,
                                       num_workers=0)

    testset = data.SSTDataset(mode='test')
    test_loader = gluon.data.DataLoader(dataset=testset,
                                        batch_size=100,
                                        batchify_fn=batcher(ctx),
                                        shuffle=False, num_workers=0)

    model = TreeLSTM(trainset.vocab_size,
                     args.x_size,
                     args.h_size,
                     trainset.num_classes,
                     args.dropout,
                     cell_type='childsum' if args.child_sum else 'nary',
                     pretrained_emb = trainset.pretrained_emb,
                     ctx=ctx)
    print(model)
    params_ex_emb =[x for x in model.collect_params().values()
                    if x.grad_req != 'null' and x.shape[0] != trainset.vocab_size]
    params_emb = list(model.embedding.collect_params().values())
    for p in params_emb:
        p.lr_mult = 0.1

    model.initialize(mx.init.Xavier(magnitude=1), ctx=ctx)
    model.hybridize()
    trainer = gluon.Trainer(model.collect_params('^(?!embedding).*$'), 'adagrad',
                            {'learning_rate': args.lr, 'wd': args.weight_decay})
    trainer_emb = gluon.Trainer(model.collect_params('^embedding.*$'), 'adagrad',
                                {'learning_rate': args.lr})

    dur = []
    L = gluon.loss.SoftmaxCrossEntropyLoss(axis=1)
    for epoch in range(args.epochs):
        t_epoch = time.time()
        for step, batch in enumerate(train_loader):
            g = batch.graph
            n = g.number_of_nodes()

            # TODO begin_states function?
            h = mx.nd.zeros((n, args.h_size), ctx=ctx)
            c = mx.nd.zeros((n, args.h_size), ctx=ctx)
            if step >= 3:
                t0 = time.time() # tik
            with mx.autograd.record():
                pred = model(batch, h, c)
                loss = L(pred, batch.label)

            loss.backward()
            trainer.step(args.batch_size)
            trainer_emb.step(args.batch_size)

            if step >= 3:
                dur.append(time.time() - t0) # tok

            if step > 0 and step % args.log_every == 0:
                pred = pred.argmax(axis=1).astype(batch.label.dtype)
                acc = (batch.label == pred).sum()
                root_ids = [i for i in range(batch.graph.number_of_nodes()) if batch.graph.out_degree(i)==0]
                root_acc = np.sum(batch.label.asnumpy()[root_ids] == pred.asnumpy()[root_ids])

                print("Epoch {:05d} | Step {:05d} | Loss {:.4f} | Acc {:.4f} | Root Acc {:.4f} | Time(s) {:.4f}".format(
                    epoch, step, loss.sum().asscalar(), 1.0*acc.asscalar()/len(batch.label), 1.0*root_acc/len(root_ids), np.mean(dur)))
        print('Epoch {:05d} training time {:.4f}s'.format(epoch, time.time() - t_epoch))

        # eval on dev set
        accs = []
        root_accs = []
        for step, batch in enumerate(dev_loader):
            g = batch.graph
            n = g.number_of_nodes()
            h = mx.nd.zeros((n, args.h_size), ctx=ctx)
            c = mx.nd.zeros((n, args.h_size), ctx=ctx)
            pred = model(batch, h, c).argmax(1).astype(batch.label.dtype)

            acc = (batch.label == pred).sum().asscalar()
            accs.append([acc, len(batch.label)])
            root_ids = [i for i in range(batch.graph.number_of_nodes()) if batch.graph.out_degree(i)==0]
            root_acc = np.sum(batch.label.asnumpy()[root_ids] == pred.asnumpy()[root_ids])
            root_accs.append([root_acc, len(root_ids)])

        dev_acc = 1.0*np.sum([x[0] for x in accs])/np.sum([x[1] for x in accs])
        dev_root_acc = 1.0*np.sum([x[0] for x in root_accs])/np.sum([x[1] for x in root_accs])
        print("Epoch {:05d} | Dev Acc {:.4f} | Root Acc {:.4f}".format(
            epoch, dev_acc, dev_root_acc))

        if dev_root_acc > best_dev_acc:
            best_dev_acc = dev_root_acc
            best_epoch = epoch
            model.save_parameters('best_{}.params'.format(args.seed))
        else:
            if best_epoch <= epoch - 10:
                break

        # lr decay
        trainer.set_learning_rate(max(1e-5, trainer.learning_rate*0.99))
        print(trainer.learning_rate)
        trainer_emb.set_learning_rate(max(1e-5, trainer_emb.learning_rate*0.99))
        print(trainer_emb.learning_rate)

    # test
    model.load_parameters('best_{}.params'.format(args.seed))
    accs = []
    root_accs = []
    for step, batch in enumerate(test_loader):
        g = batch.graph
        n = g.number_of_nodes()
        h = mx.nd.zeros((n, args.h_size), ctx=ctx)
        c = mx.nd.zeros((n, args.h_size), ctx=ctx)
        pred = model(batch, h, c).argmax(axis=1).astype(batch.label.dtype)

        acc = (batch.label == pred).sum().asscalar()
        accs.append([acc, len(batch.label)])
        root_ids = [i for i in range(batch.graph.number_of_nodes()) if batch.graph.out_degree(i)==0]
        root_acc = np.sum(batch.label.asnumpy()[root_ids] == pred.asnumpy()[root_ids])
        root_accs.append([root_acc, len(root_ids)])

    test_acc = 1.0*np.sum([x[0] for x in accs])/np.sum([x[1] for x in accs])
    test_root_acc = 1.0*np.sum([x[0] for x in root_accs])/np.sum([x[1] for x in root_accs])
    print('------------------------------------------------------------------------------------')
    print("Epoch {:05d} | Test Acc {:.4f} | Root Acc {:.4f}".format(
        best_epoch, test_acc, test_root_acc))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', type=int, default=0)
    parser.add_argument('--seed', type=int, default=41)
    parser.add_argument('--batch-size', type=int, default=256)
    parser.add_argument('--child-sum', action='store_true')
    parser.add_argument('--x-size', type=int, default=300)
    parser.add_argument('--h-size', type=int, default=150)
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--log-every', type=int, default=5)
    parser.add_argument('--lr', type=float, default=0.05)
    parser.add_argument('--weight-decay', type=float, default=1e-4)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--use-glove', action='store_true')
    args = parser.parse_args()
    print(args)
    main(args)