fastNLP/reproduction/Star_transformer/train.py

from util import get_argparser, set_gpu, set_rng_seeds, add_model_args
from datasets import load_seqtag, load_sst, load_snli, EmbedLoader, MAX_LEN
import torch.nn as nn
import torch
import numpy as np
import fastNLP as FN
from fastNLP.models.star_transformer import STSeqLabel, STSeqCls, STNLICls
from fastNLP.core.const import Const as C
import sys
sys.path.append('/remote-home/yfshao/workdir/dev_fastnlp/')


g_model_select = {
    'pos': STSeqLabel,
    'ner': STSeqLabel,
    'cls': STSeqCls,
    'nli': STNLICls,
}

g_emb_file_path = {'en': '/remote-home/yfshao/workdir/datasets/word_vector/glove.840B.300d.txt',
                   'zh': '/remote-home/yfshao/workdir/datasets/word_vector/cc.zh.300.vec'}

g_args = None
g_model_cfg = None


def get_ptb_pos():
    pos_dir = '/remote-home/yfshao/workdir/datasets/pos'
    pos_files = ['train.pos', 'dev.pos', 'test.pos', ]
    return load_seqtag(pos_dir, pos_files, [0, 1])


def get_ctb_pos():
    ctb_dir = '/remote-home/yfshao/workdir/datasets/ctb9_hy'
    files = ['train.conllx', 'dev.conllx', 'test.conllx']
    return load_seqtag(ctb_dir, files, [1, 4])


def get_conll2012_pos():
    path = '/remote-home/yfshao/workdir/datasets/ontonotes/pos'
    files = ['ontonotes-conll.train',
             'ontonotes-conll.dev',
             'ontonotes-conll.conll-2012-test']
    return load_seqtag(path, files, [0, 1])


def get_conll2012_ner():
    path = '/remote-home/yfshao/workdir/datasets/ontonotes/ner'
    files = ['bieso-ontonotes-conll-ner.train',
             'bieso-ontonotes-conll-ner.dev',
             'bieso-ontonotes-conll-ner.conll-2012-test']
    return load_seqtag(path, files, [0, 1])


def get_sst():
    path = '/remote-home/yfshao/workdir/datasets/SST'
    files = ['train.txt', 'dev.txt', 'test.txt']
    return load_sst(path, files)


def get_snli():
    path = '/remote-home/yfshao/workdir/datasets/nli-data/snli_1.0'
    files = ['snli_1.0_train.jsonl',
             'snli_1.0_dev.jsonl', 'snli_1.0_test.jsonl']
    return load_snli(path, files)


g_datasets = {
    'ptb-pos': get_ptb_pos,
    'ctb-pos': get_ctb_pos,
    'conll-pos': get_conll2012_pos,
    'conll-ner': get_conll2012_ner,
    'sst-cls': get_sst,
    'snli-nli': get_snli,
}


def load_pretrain_emb(word_v, lang='en'):
    print('loading pre-train embeddings')
    emb = EmbedLoader.fast_load_embedding(300, g_emb_file_path[lang], word_v)
    emb /= np.linalg.norm(emb, axis=1, keepdims=True)
    emb = torch.tensor(emb, dtype=torch.float32)
    print('embedding mean: {:.6}, std: {:.6}'.format(emb.mean(), emb.std()))
    emb[word_v.padding_idx].fill_(0)
    return emb


class MyCallback(FN.core.callback.Callback):
    def on_train_begin(self):
        super(MyCallback, self).on_train_begin()
        self.init_lrs = [pg['lr'] for pg in self.optimizer.param_groups]

    def on_backward_end(self):
        nn.utils.clip_grad.clip_grad_norm_(self.model.parameters(), 5.0)

    def on_step_end(self):
        warm_steps = 6000
        # learning rate warm-up & decay
        if self.step <= warm_steps:
            for lr, pg in zip(self.init_lrs, self.optimizer.param_groups):
                pg['lr'] = lr * (self.step / float(warm_steps))

        elif self.step % 3000 == 0:
            for pg in self.optimizer.param_groups:
                cur_lr = pg['lr']
                pg['lr'] = max(1e-5, cur_lr*g_args.lr_decay)



def train():
    seed = set_rng_seeds(1234)
    print('RNG SEED {}'.format(seed))
    print('loading data')
    ds_list, word_v, tag_v = g_datasets['{}-{}'.format(
        g_args.ds, g_args.task)]()
    print(ds_list[0][:2])
    embed = load_pretrain_emb(word_v, lang='zh' if g_args.ds == 'ctb' else 'en')
    g_model_cfg['num_cls'] = len(tag_v)
    print(g_model_cfg)
    g_model_cfg['init_embed'] = embed
    model = g_model_select[g_args.task.lower()](**g_model_cfg)

    def init_model(model):
        for p in model.parameters():
            if p.size(0) != len(word_v):
                nn.init.normal_(p, 0.0, 0.05)
    init_model(model)
    train_data = ds_list[0]
    dev_data = ds_list[2]
    test_data = ds_list[1]
    print(tag_v.word2idx)

    if g_args.task in ['pos', 'ner']:
        padding_idx = tag_v.padding_idx
    else:
        padding_idx = -100
    print('padding_idx ', padding_idx)
    loss = FN.CrossEntropyLoss(padding_idx=padding_idx)
    metrics = {
        'pos': (None, FN.AccuracyMetric()),
        'ner': ('f', FN.core.metrics.SpanFPreRecMetric(
            tag_vocab=tag_v, encoding_type='bmeso', ignore_labels=[''], )),
        'cls': (None, FN.AccuracyMetric()),
        'nli': (None, FN.AccuracyMetric()),
    }
    metric_key, metric = metrics[g_args.task]
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    ex_param = [x for x in model.parameters(
    ) if x.requires_grad and x.size(0) != len(word_v)]
    optim_cfg = [{'params': model.enc.embedding.parameters(), 'lr': g_args.lr*0.1},
                 {'params': ex_param, 'lr': g_args.lr, 'weight_decay': g_args.w_decay}, ]
    trainer = FN.Trainer(train_data=train_data, model=model, optimizer=torch.optim.Adam(optim_cfg), loss=loss,
                         batch_size=g_args.bsz, n_epochs=g_args.ep, print_every=10, dev_data=dev_data, metrics=metric,
                         metric_key=metric_key, validate_every=3000, save_path=g_args.log, use_tqdm=False,
                         device=device, callbacks=[MyCallback()])

    trainer.train()
    tester = FN.Tester(data=test_data, model=model, metrics=metric,
                       batch_size=128, device=device)
    tester.test()


def test():
    pass


def infer():
    pass


run_select = {
    'train': train,
    'test': test,
    'infer': infer,
}


def main():
    global g_args, g_model_cfg
    import signal

    def signal_handler(signal, frame):
        raise KeyboardInterrupt
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    parser = get_argparser()
    parser.add_argument('--task', choices=['pos', 'ner', 'cls', 'nli'])
    parser.add_argument('--mode', choices=['train', 'test', 'infer'])
    parser.add_argument('--ds', type=str)
    add_model_args(parser)
    g_args = parser.parse_args()
    print(g_args.__dict__)
    set_gpu(g_args.gpu)
    g_model_cfg = {
        'init_embed': (None, 300),
        'num_cls': None,
        'hidden_size': g_args.hidden,
        'num_layers': 4,
        'num_head': g_args.nhead,
        'head_dim': g_args.hdim,
        'max_len': MAX_LEN,
        'cls_hidden_size': 600,
        'emb_dropout': 0.3,
        'dropout': g_args.drop,
    }
    run_select[g_args.mode.lower()]()


if __name__ == '__main__':
    main()