Exemplo n.º 1
0
def train(args,model,processor):
    train_dataset = load_and_cache_examples(args, processor, data_type='train')
    train_loader = DatasetLoader(data=train_dataset, batch_size=args.batch_size,
                                 shuffle=False, seed=args.seed, sort=True,
                                 vocab = processor.vocab,label2id = args.label2id)
    parameters = [p for p in model.parameters() if p.requires_grad]
    optimizer = optim.Adam(parameters, lr=args.learning_rate)
    scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3,
                                  verbose=1, epsilon=1e-4, cooldown=0, min_lr=0, eps=1e-8)
    best_f1 = 0
    for epoch in range(1, 1 + args.epochs):
        print(f"Epoch {epoch}/{args.epochs}")
        pbar = ProgressBar(n_total=len(train_loader), desc='Training')
        train_loss = AverageMeter()
        model.train()
        assert model.training
        for step, batch in enumerate(train_loader):
            input_ids, input_mask, input_tags, input_lens = batch
            input_ids = input_ids.to(args.device)
            input_mask = input_mask.to(args.device)
            input_tags = input_tags.to(args.device)
            features, loss = model.forward_loss(input_ids, input_mask, input_lens, input_tags)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_norm)
            optimizer.step()
            optimizer.zero_grad()
            pbar(step=step, info={'loss': loss.item()})
            train_loss.update(loss.item(), n=1)
        print(" ")
        train_log = {'loss': train_loss.avg}
        if 'cuda' in str(args.device):
            torch.cuda.empty_cache()
        eval_log, class_info = evaluate(args,model,processor)
        logs = dict(train_log, **eval_log)
        show_info = f'\nEpoch: {epoch} - ' + "-".join([f' {key}: {value:.4f} ' for key, value in logs.items()])
        logger.info(show_info)
        scheduler.epoch_step(logs['eval_f1'], epoch)
        if logs['eval_f1'] > best_f1:
            logger.info(f"\nEpoch {epoch}: eval_f1 improved from {best_f1} to {logs['eval_f1']}")
            logger.info("save model to disk.")
            best_f1 = logs['eval_f1']
            if isinstance(model, nn.DataParallel):
                model_stat_dict = model.module.state_dict()
            else:
                model_stat_dict = model.state_dict()
            state = {'epoch': epoch, 'arch': args.arch, 'state_dict': model_stat_dict}
            model_path = args.output_dir / 'best-model.bin'
            torch.save(state, str(model_path))
            print("Eval Entity Score: ")
            for key, value in class_info.items():
                info = f"Subject: {key} - Acc: {value['acc']} - Recall: {value['recall']} - F1: {value['f1']}"
                logger.info(info)
Exemplo n.º 2
0
def train(args, model, processor):
    tokenizer = BertTokenizer.from_pretrained(
        './BERT_model/bert_pretrain/vocab.txt')

    train_dataset = load_and_cache_examples(args, processor, data_type='train')
    train_loader = DatasetLoader(data=train_dataset,
                                 batch_size=args.batch_size,
                                 shuffle=False,
                                 seed=args.seed,
                                 sort=True,
                                 vocab=processor.vocab,
                                 label2id=args.label2id,
                                 tokenizer=tokenizer)
    # train_loader = DatasetLoader(data=train_dataset, batch_size=args.batch_size,
    #                              shuffle=False, seed=args.seed, sort=True,
    #                              vocab=processor.vocab, label2id=args.label2id)
    parameters = [p for p in model.parameters() if p.requires_grad]
    optimizer = optim.Adam(parameters, lr=args.learning_rate)
    scheduler = ReduceLROnPlateau(optimizer,
                                  mode='max',
                                  factor=0.5,
                                  patience=3,
                                  verbose=1,
                                  epsilon=1e-4,
                                  cooldown=0,
                                  min_lr=0,
                                  eps=1e-8)

    train_metric = SeqEntityScore(args.id2label, markup=args.markup)
    best_f1 = 0
    for epoch in range(1, 1 + args.epochs):
        strat_epoch_time = time.time()
        logger.info(f"Epoch {epoch}/{args.epochs}")
        #pbar = ProgressBar(n_total=len(train_loader), desc='Training') #进度条样式
        train_loss = AverageMeter()
        model.train()
        assert model.training
        for step, batch in enumerate(train_loader):
            strat_batch_time = time.time()
            input_ids, input_mask, input_tags, input_lens = batch
            input_ids = input_ids.to(args.device)
            input_mask = input_mask.to(args.device)
            input_tags = input_tags.to(args.device)

            features, loss = model.forward_loss(input_ids, input_mask,
                                                input_lens, input_tags)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_norm)
            optimizer.step()
            optimizer.zero_grad()

            # pbar(step=step, info={'loss': loss.item()})
            train_loss.update(loss.item(), n=1)

            tags, _ = model.crf._obtain_labels(features, args.id2label,
                                               input_lens)
            input_tags = input_tags.cpu().numpy()
            target = [
                input_[:len_] for input_, len_ in zip(input_tags, input_lens)
            ]

            pre_train = train_metric.compute_train_pre(label_paths=target,
                                                       pred_paths=tags)
            logger.info(
                f'time: {time.time() - strat_batch_time:.1f}  train_loss: {loss.item():.4f}  train_pre: {pre_train:.4f}'
            )
        print(" ")
        logger.info(f'train_total_time: {time.time() - strat_epoch_time}')

        if 'cuda' in str(args.device):
            torch.cuda.empty_cache()  # 释放显存

        strat_eval_time = time.time()
        eval_f1 = evaluate(args, model, processor)

        show_info = f'eval_time: {time.time() - strat_eval_time:.1f}   train_avg_loss: {train_loss.avg:.4f}  eval_f1: {eval_f1:.4f} '
        logger.info(show_info)
        scheduler.epoch_step(eval_f1, epoch)

        if eval_f1 > best_f1:
            # Epoch 1: eval_f1 improved from 0 to 0.4023105674481821
            logger.info(
                f"\nEpoch {epoch}: eval_f1 improved from {best_f1} to {eval_f1}"
            )

            best_f1 = eval_f1

            model_stat_dict = model.state_dict()
            state = {
                'epoch': epoch,
                'arch': args.arch,
                'state_dict': model_stat_dict
            }
            model_path = args.output_dir / 'best-model.bin'
            torch.save(state, str(model_path))