BERT+BiLSTM+CRF命名实体识别实战:源码拆解与训练避坑指南

发布时间:2026/10/1 12:59:43
BERT+BiLSTM+CRF命名实体识别实战:源码拆解与训练避坑指南 简介这是一份基于PyTorch实现的命名实体识别NER完整项目核心思路是将BERT、BiLSTM与CRF三层模型串联用于识别人名、地名、组织名等实体标签为Python。资源面向NLP初学者或想快速搭建序列标注模型的开发者代码模块划分清晰从数据预处理、模型定义到训练、预测与评估都有对应脚本。压缩包共19个文件其中包含6个Python脚本如模型定义、训练主程序、CRF解码与工具函数2个Jupyter Notebook 演示数据处理和预测流程5个txt文件提供训练/测试/开发集及处理后的bio标注语料另有配置、说明文档等整体大小仅3.95MB轻量易取。目前已有936人学习浏览适合对照代码逐段理解BERTBiLSTMCRF的融合机制。通过该项目可以快速跑通NER全流程同时能基于现有代码修改数据集或模型参数开展二次实验是入门NLP序列标注的不错参考。1. 先搞清楚这个仓库到底解决什么问题做序列标注命名实体识别、分词、词性标注的应该都见过这个组合BERT 负责把上下文语义编码成向量BiLSTM 在 BERT 输出的基础上继续抽象序列特征最后 CRF 层学习标签之间的转移约束。这个 PyTorch 仓库把三件套串成了一条完整可复现的流水线从数据预处理、模型训练到推理预测全部包含适合正在学 NER 或需要拿基线模型对比效果的人。它的代码量不大但结构清晰没有把 BERT、BiLSTM、CRF 三者糊成一个黑匣子而是分文件实现改起来比较顺手。这套代码我拆过不止一次每次给别人讲 NER 的 pipeline 都会拿它当起点。下面按「源码结构 → 数据格式 → 训练流程 → 踩坑 → 预测进阶」的顺序把每个文件的职责、参数怎么调、哪些地方容易翻车一次性讲透。2. 仓库源码拆解每个文件在管线里的位置2.1 先看整体八个 Python 文件职责分得很清楚解压后能看到核心文件就这几个config.py放超参数model.py定义 BERT BiLSTM CRF 的模型结构crf.py是条件随机场的实现main.py是训练入口utils.py做数据加载和 batch 化crf_predict.py和pred.ipynb是推理脚本。我的习惯是先读config.py再读model.py因为所有关键参数都集中在 config 里读懂了它剩下代码的意图就明白一半。打开config.py大概是这样import torch class Config: # 数据路径 train_data data/example.train dev_data data/example.dev test_data data/example.test # 模型结构 bert_path bert-base-chinese # 本地下载好的模型目录或 huggingface 模型名 hidden_size 768 # BERT 输出维度base 模型固定是 768 lstm_hidden 256 # BiLSTM 隐层维度 lstm_layers 1 # BiLSTM 层数 dropout 0.5 # 全连接层前的 dropout # 训练参数 max_len 64 # 每条样本截断/补齐到该长度 batch_size 16 epochs 10 lr 3e-5 # BERT 层的学习率 lstm_lr 1e-3 # BiLSTM 和 CRF 的学习率 use_cuda torch.cuda.is_available() seed 42这里有几个参数值得注意。bert_path如果填 bert-base-chinese运行时会尝试从 HuggingFace 下载模型网络不好就需要提前把模型下载到本地再填本地路径否则大概率卡死在下载阶段。lr和lstm_lr分开设置是训练这类复合模型的常规做法——BERT 层用较小的学习率做微调下游的 BiLSTM 和 CRF 层用稍大的学习率从头学。2.2 模型是怎么搭起来的model.py的层次结构model.py定义了一个类把三个组件串成前向传播。核心结构简化后是这样import torch import torch.nn as nn from transformers import BertModel from crf import CRF class BertBiLSTMCRF(nn.Module): def __init__(self, config): super().__init__() self.bert BertModel.from_pretrained(config.bert_path) self.dropout nn.Dropout(config.dropout) self.bilstm nn.LSTM( input_sizeconfig.hidden_size, # BERT 输出的向量维度 768 hidden_sizeconfig.lstm_hidden, # 256 num_layersconfig.lstm_layers, bidirectionalTrue, # 双向 LSTM batch_firstTrue ) # 双向 LSTM 的 output 维度是 hidden_size * 2 self.fc nn.Linear(config.lstm_hidden * 2, config.num_tags) self.crf CRF(num_tagsconfig.num_tags) def forward(self, input_ids, attention_mask, labelsNone): # BERT 编码 outputs self.bert(input_ids, attention_maskattention_mask) sequence_output outputs.last_hidden_state # [batch, seq_len, 768] # BiLSTM 继续提取序列特征 lstm_out, _ self.bilstm(sequence_output) # [batch, seq_len, 512] lstm_out self.dropout(lstm_out) # 线性层映射到标签空间 logits self.fc(lstm_out) # [batch, seq_len, num_tags] if labels is not None: # CRF 计算损失负对数似然 loss -self.crf(logits, labels, maskattention_mask.bool()) return loss, logits else: # 预测时用维特比解码 decode self.crf.decode(logits, maskattention_mask.bool()) return decode, logitsattention_mask同时传给 BERT 和 CRF这是一个容易被忽略的细节。padding 的位置不参与标签预测CRF 在计算转移分数和维特比解码时必须知道哪些位置是有效 token否则会学到错误的转移概率。还要注意 BiLSTM 输出维度是hidden_size * 2。因为双向 LSTM 会拼接正向和反向两个方向的隐状态所以fc层的输入维度是 512 而不是 256。很多人在自己改代码时容易在这里维度对不上报错信息通常是size mismatch。2.3crf.py到底做了什么crf.py里的 CRF 类封装了转移矩阵、发射分数和前向计算。核心方法有三个forward函数计算配分函数decode用维特比算法求最优标签序列_score_sentence计算单条真实标签序列的得分。代码量不大但它是整个模型里最难手写的部分。一个很典型的实现片段是计算转移分数def _score_sentence(self, emissions, tags, mask): # emissions: [seq_len, batch, num_tags] # tags: [seq_len, batch] batch_size tags.shape[1] score torch.zeros(batch_size) # start 标签转移到第一个真实标签 score self.start_transitions[tags[0]] for i in range(1, seq_len): score self.transitions[tags[i - 1], tags[i]] # 最后一个标签转移到 end score self.end_transitions[tags[-1]] return score用 CRF 层的意义在于它显式建模了标签之间的依赖关系。比如在 BIO 标注体系里I-PER不能出现在O之后这个约束不是靠 BERT 学出来的而是 CRF 的转移矩阵直接编码的。如果去掉 CRF 层模型输出的标签序列大概率会出现不合法跳转。解释为什么 BERT 都已经很厉害了还非要叠 CRF这是最重要的理由。3. 数据准备和预处理BIO 标注格式与 Dataloader 的坑3.1 原始数据长什么样三列变成两列仓库自带的data/example.train、data/example.test、data/example.dev是三个小型演示数据集。打开example.train可以看到文本内容已经按 token 切好每个 token 一行和它的标签用空格或 tab 分开句子之间用空行隔开我 O 在 O 北 B-LOC 京 I-LOC 上 O 班 O 他 O 是 O 张 B-PER 三 I-PER这个格式是序列标注任务最常见的 BIO 标注格式。B-表示实体的第一个 tokenI-表示实体内部的 tokenO表示非实体。需要特别留意的是example.train只有两列而有些开源数据是「token 词性 标签」三列这类数据需要先做列裁剪。3.2data_process.ipynb做了什么仓库里有一个 Notebook 专门负责把原始文本转成训练需要的格式核心是几个转换逻辑把一句话里的 token 序列合并成完整句子同时保留每个 token 对应的 label。给 BERT tokenizer 不认识的 token 做特殊处理——中文字符通常一个字符就是一个 token但英文字母组合会被 tokenizer 拆成多个 subword此时 label 只能挂在第一个 subword 上其余 subword 统一标为X或在 loss 计算时屏蔽。输出统一格式每行token label空行分隔句子。这个环节最容易翻车的地方是「token 和 label 对不齐」。BERT 的 tokenizer 会把词语拆开比如“LSTM”有可能被拆成“LS”、“##TM”两个 token而原始 label 只给了LSTM一个标签。正确处理方式是只在第一个 subword 上保留标签后面的 subword 用X占位并在计算 loss 的时候把X对应的位置排除掉。对应的加载逻辑在utils.py里它负责把文本转成input_ids、attention_mask、labels三个张量。核心部分是这样的def build_batch(data, tokenizer, max_len): input_ids [] attention_masks [] labels [] for tokens, label_list in data: bert_tokens [] bert_labels [] for token, label in zip(tokens, label_list): # 用 BERT tokenizer 切分 token subwords tokenizer.tokenize(token) bert_tokens.extend(subwords) # 第一个 subword 保留标签其余用 -100 标记计算 loss 时忽略 bert_labels.append(label2id[label]) for _ in range(len(subwords) - 1): bert_labels.append(-100) # 截断 bert_tokens bert_tokens[:max_len - 2] bert_labels bert_labels[:max_len - 2] # 拼上 [CLS] 和 [SEP] input_ids.append(tokenizer.build_inputs_with_special_tokens(bert_tokens)) attention_masks.append([1] * len(input_ids[-1])) # 同样在标签序列前后拼上 -100 labels.append([-100] bert_labels [-100]) # padding 到 fixed length pad_len max_len - len(input_ids[-1]) input_ids[-1] [0] * pad_len attention_masks[-1] [0] * pad_len labels[-1] [-100] * pad_len return torch.tensor(input_ids), torch.tensor(attention_masks), torch.tensor(labels)逻辑说明bert_labels里的-100不是随便挑的数字PyTorch 的CrossEntropyLoss默认忽略ignore_index-100所以这些位置的预测结果不会影响梯度更新。编码后的标签序列在头尾也要拼上-100对应 BERT 的[CLS]和[SEP]位。参数说明max_len - 2是为[CLS]和[SEP]预留的两个位置。如果不用tokenizer.build_inputs_with_special_tokens而是手动[101] ids [102]记得把标签序列也手动补齐这一步做错会在训练时出现「标签长度和 logits 长度不一致」的报错。3.3 填充和注意力掩码一个容易忽视的细节attention_masks中 padding 位置填 0有效位置填 1。这个 mask 会同时传给 BERT 和 CRF。在 CRF 的解码阶段mask 用来屏蔽 padding 位置防止维特比算法把结束标签错误地转移到 padding 位上。我的建议是训练前先打印一个 batch 的数据形状确认input_ids、attention_mask、labels三个张量的维度都是[batch_size, max_len]再检查前 5 个样本的标签对齐情况。这一步只要跑一次 demo能省掉后面好几轮的排查时间。4. 环境搭建与训练全流程从零跑到第一个 epoch4.1 环境配置版本要对齐否则报错到怀疑人生这个项目的依赖主要是 PyTorch 和 Transformers。以我常用的环境为例conda create -n ner python3.8 conda activate ner pip install torch1.10.1 pip install transformers4.15.0 pip install jupyter版本选择的逻辑PyTorch 1.10.1 在 CUDA 11.3 下工作稳定Transformers 4.15.0 的BertModel返回的last_hidden_state接口和这个仓库的代码完全匹配。如果你装的是更新的 Transformers 5.xoutput.last_hidden_state依然可用但部分接口返回类型发生了变化会出现tuple index out of range之类的怪异报错排查成本很高。装完依赖后如果bert_path没有指定本地路径代码会尝试在线下载bert-base-chinese模型。这个模型大约 400MB如果网络状况不佳建议先去 HuggingFace 手动下载到本地然后把config.py里的bert_path改成对应的本地目录。4.2 跑通 demo训练脚本的执行路径数据和环境就绪后直接执行主脚本python main.py观察输出正常情况下每个 epoch 结束后会打印验证集的 F1 值。一个典型的收敛过程长这样Epoch 1/10 train loss: 12.542, val f1: 0.321 Epoch 2/10 train loss: 7.250, val f1: 0.681 Epoch 3/10 train loss: 3.610, val f1: 0.814 ...main.py的训练循环大致做了四件事加载数据、构建 batch、前向传播计算 loss、反向传播更新梯度。for epoch in range(config.epochs): model.train() total_loss 0 for input_ids, attention_mask, labels in train_loader: input_ids input_ids.to(device) attention_mask attention_mask.to(device) labels labels.to(device) optimizer.zero_grad() loss, _ model(input_ids, attention_mask, labels) loss.backward() optimizer.step() total_loss loss.item()其中优化器设置了两个参数组optimizer torch.optim.AdamW([ {params: model.bert.parameters(), lr: config.lr}, {params: model.bilstm.parameters(), lr: config.lstm_lr}, {params: model.fc.parameters(), lr: config.lstm_lr}, {params: model.crf.parameters(), lr: config.lstm_lr}, ])BERT 层用 3e-5下游层用 1e-3这种差异化学习率的做法是迁移学习的常规操作。如果下游层学习率也设成 3e-5训练速度会非常慢如果 BERT 层学习率设成 1e-3BERT 的预训练权重会被迅速破坏表现为 loss 不降反升。4.3 显存估算参数参考用max_len64、batch_size16的配置在 8GB 显存上可以跑得动bert-base-chinese。如果要加大max_len或batch_size一个估算经验是BERT base 的参数量约为 110M激活值显存占用大约按序列长度线性增长。从max_len64调到128显存需求大约翻一倍半。如果显存不够先降batch_size不要急着降max_len。max_len太短会把长句子截掉太多直接影响实体识别的召回率。4.4 模型保存与加载训练脚本通常会在每个 epoch 结束后保存模型torch.save(model.state_dict(), fcheckpoints/epoch_{epoch}_f1_{best_f1:.4f}.pt)保存的是state_dict不包含模型结构。加载时需要先实例化模型再load_state_dictmodel BertBiLSTMCRF(config) model.load_state_dict(torch.load(checkpoints/epoch_5_f1_0.812.pt)) model.to(device)因为BertModel.from_pretrained在实例化时会重新加载 BERT 权重所以加载state_dict时不要用strictFalse否则 BERT 层的权重缺失问题会被掩盖预测结果会出现「loss 很低但预测全是 O」的诡异现象。5. 避坑手册五个最容易翻车的地方5.1 标签维度对不齐Expected input batch_size to match target现象训练到一半报错Expected input batch_size ... to match target或者size mismatch。原因绝大多数情况是data_process.ipynb处理后的标签序列长度和 BERT tokenizer 切出来的 token 长度不一致。常见诱因有两个一是 token 被切分成多个 subword 时没有给后续 subword 补-100导致标签序列比 input_ids 短二是手动在原始 token 序列基础上2模拟[CLS]和[SEP]但忘了在标签序列里同样处理。解决在utils.py的build_batch函数的 return 之前加一行断言assert input_ids.shape labels.shape, fshape mismatch: {input_ids.shape} vs {labels.shape}这行代码应该保证每个样本都对齐如果跑 demo 时断言通过后面大批量训练时就不会再出这个问题。5.2 预测结果全是 O模型看起来没有任何区分能力现象训练 loss 正常下降F1 值看起来不错但实际预测句子时所有 token 都是 O。原因一个是标签映射表label2id和id2label搞反了模型输出的 ID 解析成了错误的标签另一个是加载模型时strictFalseBERT 层权重没有真正加载模型实际是随机初始化的。解决推理前先跑pred.ipynb里的 demo 验证不要直接上自己的数据。如果 demo 也全是 O打印出model.bert第 0 层第一个权重矩阵的均值和标准差和原始模型的对比数值差异大就说明 BERT 权重没加载上。5.3 中文数据分词后标签迁移到 subword 上导致信息泄漏现象验证集 F1 很高能达到 95% 以上但测试集 F1 明显偏低只有 60% 左右。原因data_process.ipynb里如果用了tokenizer做中文分词且没有处理好 subword 标签映射可能不经意间把标签信息泄漏到了 token 层面。中文按字切分通常一个汉字就是 1 个 token但如果包含英文或数字切分后就会产生 subword。解决检查数据里是否有中英混排的文本对英文和数字 token 单独处理第一个 subword 保留标签其余 subword 全部标为-100。同时把data_process.ipynb输出的训练集和验证集分开存放防止数据去重时把同一实体的不同写法泄漏到验证集。5.4 CPU 环境下训练慢到无法忍受现象use_cudaTrue但实际用的是 CPU训练一个 epoch 耗时几小时。原因config.py里use_cuda torch.cuda.is_available()会正确检测设备但很多人在main.py里把所有 tensor 移到device时写的是device torch.device(cuda)而忽略了torch.cuda.is_available()的判断。另外to(device)的操作要在每个 batch 上执行如果模型参数和数据反复在 CPU/GPU 之间拷贝也会慢。解决用config.use_cuda判断不要硬编码 cudadevice torch.device(cuda if config.use_cuda else cpu)检查是否真的用上了 GPU在训练循环前执行print(fTraining on {device})。另外确保torch.backends.cudnn.benchmark True对固定序列长度的训练能提升不少速度。5.5 预测时显存 OOM 但训练时没事现象模型跑完训练加载 checkpoint 做单条预测时反而报 CUDA OOM。原因预测时没有设置梯度关闭模型在前向传播时仍然计算了计算图导致显存占用上升。另外单条预测直接拼接batch_size1的 tensor如果max_len设得很大比如 512也会触发 OOM。解决推理前必须加torch.no_grad()with torch.no_grad(): decode, logits model(input_ids, attention_mask)如果max_len512加上大 batch 预测时仍然 OOM把预测的max_len适当降低到训练时差不多的长度通常能解决。6. 进阶用crf_predict.py做批量预测并顺手验证解码正确性6.1 跑通crf_predict.py的完整流程crf_predict.py是实现推理的脚本。它的输入是一个文本文件输出是带标注的结果。用法大致是python crf_predict.py --model_path checkpoints/epoch_5_f1_0.812.pt --input data/example.test --output pred_result.txt脚本内部要做几件事加载 tokenizer、加载模型权重、按句读取文本、执行 CRF 解码得到标签序列、把 BIO 标签还原成实体文本。6.2 验证解码正确性维特比的一个小陷阱CRF 的decode函数用的是维特比算法实现正确时预测的标签序列一定是所有可能序列里得分最高的。一个快速验证方法是把自己模型的decode结果和一个简单贪心解码每个位置取概率最高的标签做对比在转移约束较弱或标注体系简单的情况下两者结果应该高度一致。我一般会做这样一个 sanity check构造一个没有实体的句子例如我在北京上班改成我去公司上班正确输出应该全部是 O。如果输出中出现B-LOC而不带实体文本多半是模型的上下文建模出了问题或者训练数据里这个实体类别太稀疏需要检查类别分布。跑通了这一个用例就可以把自己的业务数据套进来测试了。对新数据先看模型输出的是「合法」的 BIO 序列——比如I-PER前面必须是B-PER或I-PER如果出现I-PER前面是O说明模型对这类边缘情况没有学好把这类样本挑出来单独做训练数据会很有帮助。这个仓库作为入门 NER 的基线模型非常合适。我自己跑完之后最大的感受是BERT 部分越强CRF 层能带来的提升就越小但千万别因此删掉 CRF——它是保证输出标签序列合法性的最后一道防线。从那以后我改代码凡是改动数据格式或标签映射都会强制先跑一遍 demo 数据再上全量训练。希望帮到你。本文还有配套的精品资源点击获取