PyTorch实战：BERT+BiLSTM+CRF与RoBERTa+CRF实现命名实体识别

简介面向自然语言处理课程设计与期末大作业场景一套基于Pytorch的完整命名实体识别项目源码实现了BERTBiLSTM、RobertaCRF等主流模型代码含有详细注释与配套文档说明即使新手也能轻松看懂并直接复用。压缩包共94个文件以33个Python脚本为核心覆盖模型定义、训练、评估与预测全流程另有17个JSON配置、9个npz数据文件及日志文件等辅助资源整体体积仅12.45MB便于下载部署和使用。目前该项目已有274人学习使用。资源目录按BERT-Softmax、BERT-LSTM-CRF、BERT-CRF、BiLSTM-CRF等模块清晰划分各模型实现相互独立又便于对比且包含requirements.txt环境依赖与预处理脚本下载后简单配置即可直接运行既适合作为期末大作业的高分模板也可作为命名实体识别入门进阶的实践参考项目。1. 用PyTorch做命名实体识别BERTBiLSTM与RoBERTaCRF的组合为什么值得复现常见偏见是BERT一到位、BiLSTM和CRF都可以删掉实测在中文NER任务里BERTBiLSTMCRF仍比纯BERT加Softmax高25个F1点。这个组合把上下文编码、序列投影、标签约束切成三层每一层的问题都可以单独归因这是它适合做工程项目的关键。这篇文章要讲的就是同一个PyTorch项目骨架实现基于BERTBiLSTM和RoBERTaCRF的命名实体识别模型数据预处理、模型类、训练循环和评测代码都能直接照着写。换RoBERTa只改预训练模型名拿掉BiLSTM只改forward一行剩下逻辑不动。适合要跑中文数据集、已经有Python和PyTorch基础的人想快速确认CRF解码细节的工程师可以直接跳到第4章和第5章。2. 三大组件怎么分工BERT/RoBERTa做表示、BiLSTM做投影、CRF做解码把命名实体识别当序列标注处理时整个建模路径分两大部分先把每个中文字符编码成带上下文信息的向量再从这些向量上解码出一个满足标签约束的序列。BERTBiLSTM和RoBERTaCRF是两条成型的落地路线前者适合中小数据量后者适合想尽量少改动预训练表示的团队。三个组件的信息边界很清晰BERT/RoBERTa决定“上下文向量好不好”BiLSTM在向量上做一次面向标签任务的局部融合CRF决定“整套标签序列合不合法”。下面拆开讲。2.1 BERT与RoBERTa做上下文表征前提是中文词表对齐BERT用掩码语言模型预训练输入一个句子后每个字符都被编码成带整句信息的向量。中文场景常见的预训练模型是google-bert/bert-base-chinese和hfl/chinese-roberta-wwm-ext后者即RoBERTa-wwm-extwwm是whole word masking的意思。RoBERTa相比BERT去掉了下一句预测任务并把掩码策略改成全词掩码在中文数据集上往往能多1个点左右的中文NER F1。加载模型用transformers库就可以from transformers import AutoTokenizer, AutoModel bert_name google-bert/bert-base-chinese roberta_name hfl/chinese-roberta-wwm-ext tokenizer AutoTokenizer.from_pretrained(roberta_name) pretrained AutoModel.from_pretrained(roberta_name)这里要注意AutoModel返回的对象默认只编码不包含任何标签映射文本进入模型前需要先由tokenizer把句子转成input_ids和attention_mask两个张量。第一次运行会自动下载权重文件之后from_pretrained会命中本地缓存。代码里的tokenizer和pretrained后面会作为子模块挂进大模型类。2.2 BiLSTM夹在中间到底有没有用BERT自注意力已经能双向建模所以BiLSTM不是必须但实践中它往往带来三个好处。第一实体通常是连续出现的自然片段LSTM的递归结构对局部连续性有天然归纳偏置在小数据集上这种偏置比密集注意力更抗过拟合。第二BiLSTM可以把BERT输出的768维或1024维压缩到256维后面再接CRF时参数规模小很多。第三从调试角度看单独观察LSTM的输出比直接看BERT每个head的注意力更容易定位“句子长了之后实体尾部丢了”的问题。一个常见的配置是单向hidden设为128双向拼接成256self.lstm nn.LSTM( hidden_size, hidden_size128, bidirectionalTrue, batch_firstTrue, )batch_firstTrue让进出模型的张量形状统一成(batch, seq, hidden)PyTorch的LSTM默认是(seq, batch, hidden)很多维度错位都出在这个参数上。2.3 CRF负责全局标签转移约束CRF做的事情不是给每个字单独打分而是直接在标签序列层面上计算概率。它比Softmax多学了一个(num_tags, num_tags)的转移矩阵矩阵里存放“上一个标签是A、当前标签是B”的得分。2.3.1 转移矩阵训练的是什么例如B-LOC后面可以出现I-LOC但O后面通常不该直接接I-PERB-PER后面接B-LOC虽然不常见在矩阵里仍会被赋予一个小分数。模型训练时CRF层输出整句标签序列的负对数似然梯度会同时更新发射得分和转移得分。用torchcrf验证这个层只需要学习一小段代码import torch import torchcrf num_tags 7 crf torchcrf.CRF(num_tags, batch_firstTrue) emissions torch.randn(2, 5, num_tags) mask torch.ones(2, 5, dtypetorch.bool) decoded crf.decode(emissions, maskmask)emissions的形状是(batch, seq_len, num_tags)这是CRF层的前提条件。mask标记哪些位置是有效tokenpadding部分会被排出计算。上面代码里的decoded是两个长度各为5的标签序列不是概率值。2.3.2 Viterbi解码保证序列整体最优解码时如果每个位置独立取最大概率标签很容易出现“B-PER之后直接跳I-LOC”这类非法序列。Viterbi算法在解码过程中保留所有转移路径最后输出整句得分最高的路径计算代价是O(seq_len * num_tags^2)。当标签数量小于20时这个开销可以忽略。解码方式约束能力典型错误每个位置独立argmax无B-PER后直接I-LOCCRF Viterbi转移矩阵约束可避免上下文非法跳跃NER任务的最终评测又按实体span计算边界错一个字整体就算失败所以CRF带来的边界修正对F1影响很大。3. 准备NER环境与数据PyTorch安装、BIO标注和标签对齐一个NER项目能不能复现一半取决于训练指标另一半取决于数据预处理是否严谨。很多新手卡在“模型能跑但指标只有60”问题常常不来自模型而是标签对齐错了。3.1 创建PyTorch基础环境并安装依赖常见做法是先用conda建独立的Python环境再装PyTorch基础框架。命名实体识别需要transformers做预训练模型加载、seqeval做实体级评估、torchcrf做CRF层。conda create -n ner python3.9 -y conda activate ner pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cu118 pip install transformers seqeval tqdm torchcrf第一行创建名为ner的独立环境避免污染系统其他项目。第二行固定torch 2.1.0原因不是越新越好而是和transformers 4.3x系列的API兼容性最稳。如果你的机器没有NVIDIA显卡或CUDA驱动版本比较乱直接装CPU版pip install torch2.1.0CPU版慢但用来跑单条样本验证代码逻辑、查看张量shape完全够。我的习惯是先在CPU上把一个batch跑通再换GPU做完整训练这样能避开不少环境错误。3.2 中文NER数据加载与BIO标签格式中文NER数据通常是一个句子配一组等长的BIO标签实体类别常用PER人名、LOC地点、ORG组织。合成一条示例数据text: 张三在北京的阿里巴巴工作 labels: B-PER O O B-LOC O B-ORG O O处理时把标签存成字符串列表即可。JSON行格式常见如下{text: 张三在北京的阿里巴巴工作, labels: [B-PER, O, O, B-LOC, O, B-ORG, O, O]}注意不要直接把标签做独热编码CRF层需要的是整数id序列。建立标签映射label2id { O: 0, B-PER: 1, I-PER: 2, B-LOC: 3, I-LOC: 4, B-ORG: 5, I-ORG: 6, } id2label {v: k for k, v in label2id.items()} num_tags len(label2id)固定标签顺序后训练和推理必须共用同一份id2label否则评估阶段边界会乱。3.3 标签对齐把字符BIO映射成token标签中文大部分字符在BERT词表里是单字token但英文、数字、符号会被拆成子词。如果一句“OpenAI在旧金山”被tokenizer拆成Open和AI两个token而原始BIO标签是按单字符顺序给的直接zip就会错位。3.3.1 用逐字符tokenize避免常见错位一种稳妥做法是先按字符切分再交给tokenizer转成subworddef char_tokenize(char): return tokenizer.tokenize(char) def convert_example(text, label_chars, max_len128): token_ids [] label_ids [] for char, lab in zip(text, label_chars): subtokens tokenizer.tokenize(char) if not subtokens: continue for st in subtokens: token_ids.append(tokenizer.convert_tokens_to_ids(st)) label_ids.append(label2id[lab])逐字符处理的好处是任何中文字符最多被拆成12个子词标签始终跟随原始字符展开不会因为分词边界错位。3.3.2 添加CLS/SEP并统一paddingBERT类模型需要在句首加[CLS]、句尾加[SEP]这两个位置的标签要跟着处理if len(token_ids) max_len - 2: token_ids token_ids[:max_len - 2] label_ids label_ids[:max_len - 2] token_ids [tokenizer.cls_token_id] token_ids [tokenizer.sep_token_id] label_ids [0] label_ids [0] attention_mask [1] * len(token_ids)max_len - 2是因为CLS和SEP占两个位置。label_ids首尾填0即O标签因为CRF层在计算时会用attention_mask把padding位置排除只要有效位置标签正确就不影响训练。DataLoader批处理时不同句子长度不同需要padding到一致def collate_fn(batch): input_ids, label_ids, masks [], [], [] max_len max(len(x[0]) for x in batch) for ids, labels, mask in batch: ids ids [tokenizer.pad_token_id] * (max_len - len(ids)) labels labels [0] * (max_len - len(labels)) masks mask [0] * (max_len - len(mask)) input_ids.append(ids) label_ids.append(labels) masks.append(masks) return ( torch.tensor(input_ids), torch.tensor(masks, dtypetorch.bool), torch.tensor(label_ids, dtypetorch.long), )这里masks用布尔类型后面CRF层直接接收attention_mask不需要再.bool()。padding标签全填0但由于mask为False不会被计入损失。4. 用PyTorch实现BERTBiLSTMCRF与RoBERTaCRF模型模型实现是整个项目源码的核心。我把两条路线放在同一套代码框架里分别对应标题里的“BERTBiLSTM”和“RoBERTaCRF”。4.1 完整的BERTBiLSTMCRF模型类import torch from torch import nn import torchcrf from transformers import AutoModel class BertBiLstmCrf(nn.Module): def __init__(self, bert_name, num_tags, lstm_hidden256, dropout0.1): super().__init__() self.bert AutoModel.from_pretrained(bert_name) self.lstm nn.LSTM( self.bert.config.hidden_size, hidden_sizelstm_hidden // 2, bidirectionalTrue, batch_firstTrue, ) self.dropout nn.Dropout(dropout) self.classifier nn.Linear(lstm_hidden, num_tags) self.crf torchcrf.CRF(num_tags, batch_firstTrue) def forward(self, input_ids, attention_mask, labelsNone): mask attention_mask.bool() bert_output self.bert( input_idsinput_ids, attention_maskattention_mask, )[0] lstm_out, _ self.lstm(bert_output) emissions self.classifier(self.dropout(lstm_out)) if labels is not None: loss -self.crf( emissions, labels, maskmask, reductionmean, ) return loss return self.crf.decode(emissions, maskmask)bert_output的形状是(batch, seq, hidden)其中hidden在bert-base-chinese里是768。LSTM的hidden_size设置成lstm_hidden // 2双向拼接后正好是256维。classifier把256维映射到标签数得到每个token在每个标签上的得分也就是CRF的发射得分。torchcrf.CRF接收发射得分和标签序列返回负的对数似然所以forward里加负号变成损失。decode返回的是整数标签序列列表不是概率值。参数调节建议参数推荐值说明lstm_hidden256过大容易过拟合过小表达力不足dropout0.1预训练模型微调时适合偏小num_layers1单层双向LSTM通常优于多层reductionmean按batch内所有有效token取平均4.2 换成RoBERTaCRF的快速实现RoBERTa和BERT在transformers接口层面完全一致。如果不想经过BiLSTM直接让预训练模型输出后接CRFclass RobertaCrf(nn.Module): def __init__(self, roberta_name, num_tags, dropout0.2): super().__init__() self.roberta AutoModel.from_pretrained(roberta_name) self.dropout nn.Dropout(dropout) self.classifier nn.Linear( self.roberta.config.hidden_size, num_tags, ) self.crf torchcrf.CRF(num_tags, batch_firstTrue) def forward(self, input_ids, attention_mask, labelsNone): mask attention_mask.bool() roberta_output self.roberta( input_idsinput_ids, attention_maskattention_mask, )[0] emissions self.classifier(self.dropout(roberta_output)) if labels is not None: return -self.crf( emissions, labels, maskmask, reductionmean, ) return self.crf.decode(emissions, maskmask)这个版本的参数数量比BERTBiLSTM少训练速度更快。如果你的预训练权重本身已经是在中文语料上充分训练过的RoBERTa直接接CRF往往就能打平甚至超过BERTBiLSTMCRF。是否需要BiLSTM可以从验证集F1来判断不要凭感觉保留。4.2.1 两条路线的选择逻辑数据量在几千句级别时建议用BERTBiLSTMCRF因为BiLSTM能补充局部连续性归纳偏置。数据量到几万句且算力充足时优先试RoBERTaCRF减少一个可调模块鲁棒性更好。这两套模型共用同一个数据加载器、同一个id2label映射切换成本只需要改一行模型类名。4.3 前向跑一次先检查张量形状写完模型类不要急着训练先构造随机输入验证形状model BertBiLstmCrf(google-bert/bert-base-chinese, num_tags7) ids torch.randint(0, 2000, (2, 16)) mask torch.ones(2, 16, dtypetorch.bool) labels torch.randint(0, 7, (2, 16)) loss model(ids, mask, labels) print(loss) decoded model(ids, mask) print(len(decoded), decoded[0] if len(decoded) 0 else empty)如果loss能正常返回一个正数说明前向传播没有问题。直接调用model(ids, mask)时如果所有句子在mask后的有效长度不一致decode返回的每条序列长度会不同这是CRF层正常行为不是bug。4.4 CRF与逐个token Softmax的核心差别4.4.1 模型结构上的差异Softmax版模型在classifier之后直接算交叉熵完全无视相邻标签关系。CRF版在classifier之后多学了一张转移矩阵。这个矩阵在训练中会为“O后面接B-PER”给出较高分数同时压低“O后面直接接I-PER”的分数于是边界错误在解码阶段就能被抑制。4.4.2 损失函数上的差异Softmax的交叉熵对每个位置独立求损失CRF还会加上整个序列的归一化项。你可以把CRF理解成对整条路径做Softmax也就意味着它天然适合“实体边界完整性”这类全局指标。5. 训练循环、超参数配置和常见PyTorch坑模型搭好后训练阶段的参数配置直接决定指标能到多少。这里给出可复制的训练循环以及几个实际会碰到的坑。5.1 参数分组和warmup调度BERT类模型微调最忌讳一个学习率走天下。常见做法是把预训练参数和新增参数分开设置学习率from transformers import get_linear_schedule_with_warmup optimizer torch.optim.AdamW( [ {params: model.bert.parameters(), lr: 3e-5}, {params: model.lstm.parameters(), lr: 3e-4}, {params: model.classifier.parameters(), lr: 3e-4}, {params: model.crf.parameters(), lr: 3e-4}, ], weight_decay0.01, ) total_steps len(train_loader) * epochs scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepsint(total_steps * 0.1), num_training_stepstotal_steps, )这里BERT参数学习率3e-5因为预训练权重只需微调新增的LSTM、classifier、CRF用3e-4收敛更快。num_warmup_steps取总步数的10%让训练初期更稳。5.2 一个完整的训练batch逻辑from torch.nn.utils import clip_grad_norm_ from tqdm import tqdm def train_one_epoch(model, dataloader, optimizer, scheduler, device): model.train() total_loss 0 for batch in tqdm(dataloader): input_ids, attention_mask, labels [x.to(device) for x in batch] optimizer.zero_grad() loss model(input_ids, attention_mask, labels) loss.backward() clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() scheduler.step() total_loss loss.item() return total_loss / len(dataloader)clip_grad_norm_把梯度范数限制在5.0防止LSTM在长句子上梯度爆炸。optimizer.zero_grad()必须在loss.backward()前执行否则梯度会在batch之间累加。scheduler.step()每个batch更新一次不是每个epoch更新一次。5.3 实体级指标用seqeval别自己数标签NER评估不能只看token准确率因为“B-PER”和“I-PER”边界错一位实体就算错。common做法是使用seqeval库计算实体级别的precision、recall、F1from seqeval.metrics import classification_report true_labels [ [B-PER, I-PER, O], [B-LOC, O, B-ORG], ] pred_labels [ [B-PER, I-PER, O], [B-LOC, O, O], ] print(classification_report(true_labels, pred_labels, digits4))classification_report会把连续的一组B/I标签合并成一个实体然后按完整实体匹配来算分。上面的示例中第二条样本的ORG完全没预测出来F1会明显降低。你需要确保传入的是label字符串而不是整数id。5.4 训练过程中真正会碰到的坑5.4.1 Windows下c10.dll初始化失败训练前如果看到OSError: [WinError 1114]指向torch/lib/c10.dll这个错误通常来自torch安装时混用了不同来源的包。比如先用conda装了一个CPU版本又用pip装GPU版本。解决方法是把环境里所有torch相关包卸掉重新安装与CUDA匹配的版本pip uninstall -y torch torchvision pip cache purge pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cu118如果机器没有独立显卡装CPU版即可。5.4.2 CRF解码结果比输入短当attention_mask中有False时torchcrf.decode会跳过padding位置返回的每条序列长度等于该样本有效token数。这会导致DataLoader里一个batch的解码结果无法直接拼成二维张量。评估时统一按原始长度处理for item, mask in zip(decoded, attention_mask): seq_len mask.sum().item() padded_pred item [0] * (seq_len - len(item))5.4.3 标签对齐错位但损失很低如果训练集损失下降正常验证集实体F1却很低最常见的故障是对齐错位。例如先把整句交给tokenizer再用原始字符标签直接zip遇到BERT把英文拆成多个subword时标签序列少一位之后所有标签全部错位。这类错位用单条样本的推理可视化最容易发现。6. 推理与字符标签映射让CRF预测结果落到原文上模型训练完最后一步是把它接到真实业务输入上。这里要解决两件事模型输出的是subword token标签需要映射回原文字符CRF会跳掉padding位置需要对齐offset。6.1 单条句子推理函数def predict_one(model, tokenizer, text, max_len128): enc tokenizer( text, max_lengthmax_len, truncationTrue, return_offsets_mappingTrue, ) input_ids torch.tensor([enc[input_ids]]) attention_mask torch.tensor([enc[attention_mask]], dtypetorch.bool) model.eval() with torch.no_grad(): decoded model(input_ids, attention_mask)[0] return enc, decodedreturn_offsets_mappingTrue会返回每个token在原字符串里的起止位置这个信息是映射回原文字符的关键。model(input_ids, attention_mask)触发的是forward里不带labels的分支返回CRF解码结果。6.2 把token标签映射到每个中文汉字def merge_to_char(text, enc, decoded, id2label): char_labels [O] * len(text) for token_label_id, offset in zip(decoded, enc[offset_mapping]): start, end offset if start end: continue label id2label[token_label_id] if label.startswith(B-) or label.startswith(I-): for pos in range(start, end): if pos len(text): char_labels[pos] label return list(zip(text, char_labels))当某个token的start end时说明是CLS、SEP或padding直接跳过。恢复出的char_labels与原始文本长度一致后续输出实体span时就可以直接按连续标签切分。6.3 部署时的两个边界情况第一如果业务文本里中文和英文混排比如“iPhone在北京发布”分词器会把iPhone拆成多个subword。推理端的offset_mapping能正确映射但训练端必须用同一套对齐逻辑否则一对多关系会在训练时被破坏。第二CRF解码速度和标签数、序列长度相关实时服务场景中建议把序列长度限制在128不要无脑放开到512。另外如果你同时发布BERT与RoBERTa两个版本模型需要彻底共用同一份id2label.json。两个模型各自建标签表很容易出现“模型A的B-PER是1、模型B的B-PER是3”这类隐患切换模型时边界全部解析失败。部署前把这个映射文件单独测试一遍比反复调阈值更有价值。本文还有配套的精品资源点击获取