基于 Bert 实现 NER 任务 – 模型训练

我们使用 Bert 的 bert-base-chinese 为基础模型，在我们自己的数据集上进行微调来实现 NER 任务。我们这里使用的是 transformers 库中提供的 BertForTokenClassification 类来进行微调，该 head 非常适合字词粒度的分类任务。

import datasets
from transformers import BertModel
from transformers import BertConfig
from transformers import BertTokenizer
from datasets import load_from_disk
from transformers import BertForMaskedLM
from transformers import BertForTokenClassification
import torch
from torch.nn.utils.rnn import pad_sequence
import torch.optim as optim
import math
import numpy as np

由于数据集有 4 万+，训练时间较长，我们这里我们使用累计多个 step 的损失之后，再更新参数，训练效率得到明显提升。训练的轮数为 40，也可以改成其他值。存储模型时，从 11 epoch 开始，每一个 epoch 结束就存储一次。

另外，由于我们希望模型对象帮我们去自动计算损失，所以，在调用 model.forward 函数时，可以传递 labels 参数，但要注意的是，我们传递的 labels 参数必须为 tensor 类型，它要求每个数据长度必须一样。所以，可以使用 pad_sequence 对 labels 填充下，模型内部会使用 attention_mask 屏蔽填充这部分标签值的。

受限于显存大小，我这里最高只能设置 batch_size 为 4，没迭代 8 次，即: 累计 8 次梯度进行一次参数更新。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def train():

    # 读取训练数据
    train_data = load_from_disk('data/03-train')['train']

    # 初始化分词器
    tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
    # 初始化模型
    model = BertForTokenClassification.from_pretrained('bert-base-chinese',
                                                       num_labels=7)
    model.to(device)
    # 优化器
    optimizer = optim.AdamW(model.parameters(), lr=3e-5)
    # 训练批次
    batch_size = 4
    # 训练轮数
    epoch_num = 40
    # 累计梯度
    accumulate_step = 0

    total_steps = int(math.ceil(len(train_data) / batch_size))


    def start_train(data_inputs, data_label_ids):

        data_inputs = tokenizer.batch_encode_plus(data_inputs,
                                                  add_special_tokens=False,
                                                  padding='longest',
                                                  return_tensors='pt')

        data_inputs['input_ids'] = data_inputs['input_ids'].to(device)
        data_inputs['attention_mask'] = data_inputs['attention_mask'].to(device)
        data_inputs['token_type_ids'] = data_inputs['token_type_ids'].to(device)

        # 标签转换张量
        data_labels = []
        for labels in data_label_ids:
            data_labels.append(torch.tensor(labels, dtype=torch.int64, device=device))

        labels = pad_sequence(data_labels, batch_first=True, padding_value=-1)

        # 模型计算
        outputs = model(**data_inputs, labels=labels)

        # 反向传播
        loss = outputs.loss
        loss.backward()

        nonlocal accumulate_step
        accumulate_step += 1

        if accumulate_step % 8 == 0 or accumulate_step == total_steps:
            # 参数更新
            optimizer.step()
            # 梯度清零
            optimizer.zero_grad()

        nonlocal epoch_loss
        epoch_loss += loss.item()

    for epoch_idx in range(epoch_num):
        epoch_loss = 0.0
        train_data.map(start_train,
                       batched=True,
                       batch_size=batch_size,
                       input_columns=['data_inputs', 'data_label_ids'])

        print('loss: %.5f' % epoch_loss)
        if epoch_idx > 10:
            model.save_pretrained('data/ner-model-%d' % epoch_idx)

训练结束后得到的模型文件如下：

ner-model-12  ner-model-17  ner-model-22  ner-model-27  ner-model-32  ner-model-37
ner-model-13  ner-model-18  ner-model-23  ner-model-28  ner-model-33  ner-model-38
ner-model-14  ner-model-19  ner-model-24  ner-model-29  ner-model-34  ner-model-39
ner-model-15  ner-model-20  ner-model-25  ner-model-30  ner-model-35
ner-model-11  ner-model-16  ner-model-21  ner-model-26  ner-model-31  ner-model-36

每个模型目录下的内容为：

ner-model-11
├── config.json
└── pytorch_model.bin

所有的模型链接：https://www.aliyundrive.com/s/yNVU6FnLrhf 提取码: x5j5