基于BERT-BiLSTM-CRF的中文命名实体识别模型训练代码

import torch import torch.nn as nn import torch.optim as optim from seqeval.metrics import classification_report, f1_score from torch.utils.data import DataLoader from tqdm import tqdm from copy import deepcopy import warnings import matplotlib.pyplot as plt from transformers import logging from dataset import LABEL, Data_set, idx2tag, padding from model.bert_bilstm_crf import BertBilstmCRF

logging.set_verbosity_warning() warnings.filterwarnings('ignore', category=DeprecationWarning) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def train(model, train_iter, optimizer, loss_fn): ''' 训练函数 ''' model.train() train_l_sum, c = 0., 0 for i, (x, y) in enumerate(train_iter): x, y = x.to(device), y.to(device) y_hat = model(x) loss = loss_fn(y_hat, y) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() train_l_sum += loss.item() c += 1

    if (i + 1) % 15 == 0:
        print(f"step: {i + 1}, cur train loss: {loss.item():.6f}")

return train_l_sum / c

def evaluate(model, eval_iter, report=False): ''' 评估模型 ''' model.eval() y_true, y_pred = [], [] with torch.no_grad(): for x, y in tqdm(eval_iter): masks = (x != 0).numpy() x = x.to(device) logits = model(x) # 将每个句子padding的去除 y_hat = model.decode(logits, masks).detach().cpu().numpy() y = y.numpy() for sentence, pred, mask in zip(y, y_hat, masks): y_true.append([ idx2tag[l] for l, m in zip(sentence, mask) if m ]) y_pred.append([ idx2tag[l] for l, m in zip(pred, mask) if m ]) eval_f1 = f1_score(y_pred=y_pred, y_true=y_true) if report: print("classification report:") print(classification_report(y_true=y_true, y_pred=y_pred))

return eval_f1

def main(args): # 加载数据集 train_set = Data_set("BIOES_data/processed/all/train.txt", args['pretrained_path']) dev_set = Data_set("BIOES_data/processed/all/dev.txt", args['pretrained_path']) test_set = Data_set("BIOES_data/processed/all/test.txt", args['pretrained_path']) train_iter = DataLoader( dataset=train_set, batch_size=args['batch_size'], collate_fn=padding, num_workers=4, pin_memory=True ) dev_iter = DataLoader( dataset=dev_set, batch_size=args['batch_size'], collate_fn=padding, num_workers=4, pin_memory=True ) test_iter = DataLoader( dataset=test_set, batch_size=args['batch_size'], collate_fn=padding, num_workers=4, pin_memory=True ) # 加载模型 model = BertBilstmCRF( output_size=args['output_size'], hidden_size=args['hidden_size'], drop_prob=args['drop_prob'], embed_size=args['embed_size'], num_layers=args['num_layers'], pretrained_path=args['pretrained_path'] ) model.to(device) loss_fn = model.neg_log_likelihood optimizer = optim.Adam(params=model.parameters(), lr=args['lr'])

best_model, best_f1 = None, 0.
train_losses, eval_f1s = [], []
for e in range(args['epochs']):
    train_loss = train(model, train_iter, optimizer, loss_fn)
    eval_f1 = evaluate(model, dev_iter)
    train_losses.append(train_loss)
    eval_f1s.append(eval_f1)
    if best_f1 < eval_f1:
        best_f1 = eval_f1
        best_model = deepcopy(model)
    print("Epoch: {}, train_loss: {:.6f}  eval_f1: {:.6f}".format(
        e + 1, train_loss, eval_f1
    ))

evaluate(best_model, test_iter, report=True)
x_ticks = list(range(1, args['epochs'] + 1, 10))
plt.figure(figsize=(8, 6))
plt.subplot(1, 2, 1)
plt.title("Train Set Loss")
plt.plot(list(range(args['epochs'])), train_losses)
plt.xticks(x_ticks)
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.subplot(1, 2, 2)
plt.title("Valid Set F1-Score")
plt.plot(list(range(args['epochs'])), eval_f1s)
plt.xticks(x_ticks)
plt.xlabel("Epoch")
plt.ylabel("f1 score")
plt.savefig("result/outcome.png")
plt.show()

if name == "main": params = { 'pretrained_path': "bert-base-chinese-ws", "lr": 0.001, "batch_size": 64, "epochs": 100, "output_size": len(LABEL), "embed_size": 768, "hidden_size": 256, "num_layers": 2, "drop_prob": 0.5 } print(params) main(params)