Upload New File

NLU/train.py

+import argparse
+import torch
+import os
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--bert_path', help='config file', default='/home/data/tmp/bert-base-chinese')
+parser.add_argument('--save_path', help='path to save checkpoints', default='/home/data/tmp/NLP_Course/Joint_NLU/train')
+parser.add_argument('--train_file', help='training data', default='/home/data/tmp/NLP_Course/Joint_NLU/data/train.tsv')
+parser.add_argument('--valid_file', help='valid data', default='/home/data/tmp/NLP_Course/Joint_NLU/data/test.tsv')
+parser.add_argument('--intent_label_vocab', help='training file', default='/home/data/tmp/NLP_Course/Joint_NLU/data/cls_vocab')
+parser.add_argument('--slot_label_vocab', help='training file', default='/home/data/tmp/NLP_Course/Joint_NLU/data/slot_vocab')
+
+parser.add_argument("--local_rank", help='used for distributed training', type=int, default=-1)
+parser.add_argument('--lr', type=float, default=8e-6)
+parser.add_argument('--lr_warmup', type=float, default=200)
+parser.add_argument('--bs', type=int, default=30)
+parser.add_argument('--batch_split', type=int, default=1)
+parser.add_argument('--eval_steps', type=int, default=40)
+parser.add_argument('--n_epochs', type=int, default=30)
+parser.add_argument('--max_length', type=int, default=90)
+parser.add_argument('--seed', type=int, default=123)
+parser.add_argument('--n_jobs', type=int, default=1, help='num of workers to process data')
+
+parser.add_argument('--gpu', help='which gpu to use', type=str, default='3')
+
+args = parser.parse_args()
+
+os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+
+from transformers import BertConfig, BertTokenizer, AdamW
+from NLU_model import NLUModule
+import dataset
+import utils
+import traceback
+from trainer import Trainer
+from torch.nn.parallel import DistributedDataParallel
+
+train_path = os.path.join(args.save_path, 'train')
+log_path = os.path.join(args.save_path, 'log')
+
+def save_func(epoch, device):
+    filename = utils.get_ckpt_filename('model', epoch)
+    torch.save(trainer.state_dict(), os.path.join(train_path, filename))
+
+try:
+    if args.local_rank == -1 or args.local_rank == 0:
+        if not os.path.isdir(args.save_path):
+            os.makedirs(args.save_path)
+    while not os.path.isdir(args.save_path):
+        pass
+    logger = utils.get_logger(os.path.join(args.save_path, 'train.log'))
+
+    # Setup logging and save folder
+    if args.local_rank == -1 or args.local_rank == 0:
+        for path in [train_path, log_path]:
+            if not os.path.isdir(path):
+                logger.info('cannot find {}, mkdiring'.format(path))
+                os.makedirs(path)
+
+        for i in vars(args):
+            logger.info('{}: {}'.format(i, getattr(args, i)))
+
+    distributed = (args.local_rank != -1)
+    if distributed:
+        # Setup distributed training
+        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(args.seed)
+    else:
+        device = torch.device("cuda", 0)
+
+    tokz = BertTokenizer.from_pretrained(args.bert_path)
+    _, intent2index, _ = utils.load_vocab(args.intent_label_vocab)
+    _, slot2index, _ = utils.load_vocab(args.slot_label_vocab)
+    train_dataset = dataset.NLUDataset([args.train_file], tokz, intent2index, slot2index, logger, max_lengths=args.max_length)
+    valid_dataset = dataset.NLUDataset([args.valid_file], tokz, intent2index, slot2index, logger, max_lengths=args.max_length)
+
+    logger.info('Building models, rank {}'.format(args.local_rank))
+    bert_config = BertConfig.from_pretrained(args.bert_path)
+    bert_config.num_intent_labels = len(intent2index)
+    bert_config.num_slot_labels = len(slot2index)
+    model = NLUModule.from_pretrained(args.bert_path, config=bert_config).to(device)
+
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
+
+    trainer = Trainer(args, model, tokz, train_dataset, valid_dataset, log_path, logger, device, distributed=distributed)
+
+    start_epoch = 0
+    if args.local_rank in [-1, 0]:
+        trainer.train(start_epoch, args.n_epochs, after_epoch_funcs=[save_func])
+    else:
+        trainer.train(start_epoch, args.n_epochs)
+
+except:
+    logger.error(traceback.format_exc())