简介基于PyTorch与BERT的实体关系抽取完整工程代码包面向自然语言处理初学者与知识图谱构建者适用于中文关系抽取任务。项目中已内置DUIE数据集加载逻辑支持通过配置文件替换为自定义数据可快速跑通训练、测试与推理流程基于PyTorch 1.5.1与Transformers 2.5.1环境测试通过。压缩包共20个文件包括5个Python源码、5个编译缓存、5张结果展示图、2个Shell启动脚本及运行日志整体仅2.45MB轻量易用Python脚本覆盖数据加载、模型定义、训练、测试与演示等关键环节Shell脚本可一键启动训练和测试另附可视化图片便于直观查看效果整体目录结构清晰、模块职责分明便于二次开发已有248人学习下载。阅读源码后可掌握BERT序列标注思路、关系分类训练细节、自定义数据接入方法以及Shell脚本执行流程适合作为入门知识图谱关系抽取的参考项目利用该方案可快速复现关系抽取实验为后续知识图谱构建打下基础。1. 实体关系抽取为什么要落在 BERT 上实体关系抽取是知识图谱构建里最费人力的环节。早期做法靠特征模板写正则匹配“位于”“出生于”这类触发词换领域就失效后来用 BiLSTM-CRF 做序列标注但关系类型一多标签体系爆炸而且无法处理跨长距离的实体依赖。用 BERT 做这件事的思路不同把关系分类变成一句“句子 头实体 尾实体”的语义匹配问题让预训练模型去判断中间有没有关系、是什么关系。DUIE 是百度信息抽取大赛的数据集里面有 schema、spo_list正好拿来做关系抽取的训练与评测。这个项目的主线是用 DUIE 数据构造关系分类样本用 bert-base-chinese 做编码接一个线性分类头输出关系类型。对 5 年以上的工程师来说值得关注的是数据对齐和采样策略而不是模型本身因为 BERT 部分已经很成熟坑大多在教学流程之外。2. 数据准备与 loader 设计从 DUIE 原始结构到训练样本2.1 DUIE 的原始格式与本次任务的差异DUIE 官方数据是 JSON 行每条包含text和spo_list。spo_list里每一项是predicate、subject、object以及它们在文本中的偏移量。下面的片段是原始格式的典型样子{ text: 《红楼梦》的作者是曹雪芹曹雪芹名霑字梦阮。, spo_list: [ { predicate: 作者, subject: 红楼梦, object: 曹雪芹, subject_offset: [0, 4], object_offset: [9, 12] } ] }注意本项目里的train.json就是这种结构。但我们要做的是关系分类不是序列标注所以需要从spo_list中抽出正样本再把不含该关系的实体对构造为负样本。关系分类的输入通常是句子 头实体 尾实体输出是关系类型。如果直接沿用原始 JSON模型会遇到两个问题一条文本里可能有多对实体、多个关系直接整句分类无法确定“谁和谁”的关系。因此 loader 的核心职责是样本构造和 token 对齐。2.2 构造关系分类样本的策略常见做法是遍历spo_list中的每个三元组把subject当成头实体object当成尾实体predicate当作标签与此同时随机抽取文本中其他实体组合作为负样本。这个项目里 loader.py 也是这么处理的具体分三步2.2.1 读取与字段解析先按行加载 JSON解析出text与spo_list。对每条 SPO生成一条正样本(text, subject, object, predicate_id)。如果某条文本同时有多个 SPO就生成多条样本保证每个三元组都能被训练到。2.2.2 负样本采样正样本的数量通常远少于错配的实体对数量。为了模型学会区分“有关系”与“无关系”需要负样本。我一般用的采样方式是固定一条文本随机选择一个出现在该文本中的实体作为头再随机选择另一个实体作为尾如果这个组合不在spo_list中就标记为“无关系”类别比如predicate为NEG。如果采样后发现正负比例超过 1:3就减少负采样数量避免模型偏向多数类。2.2.3 文本拼接与 tokenizer 编码BERT 不能直接吃原始三元组需要按固定模板拼成一个序列。项目里用的是“[CLS] 文本 [SEP] 头实体 [SEP] 尾实体 [SEP]”的格式。用tokenizer.encode_plus可以一次拿到input_ids、token_type_ids和attention_maskfrom transformers import BertTokenizer tokenizer BertTokenizer.from_pretrained(bert-base-chinese) def build_one_sample(text, subject, object_, label): encoded tokenizer.encode_plus( text, subject, object_, add_special_tokensTrue, max_length128, paddingmax_length, truncationonly_first, # 优先截断文本 return_tensorspt ) return { input_ids: encoded[input_ids].squeeze(0), token_type_ids: encoded[token_type_ids].squeeze(0), attention_mask: encoded[attention_mask].squeeze(0), label: label }注意encode_plus传三个参数时BERT 会把第二、第三个参数当作text_pair并用两个或多个[SEP]分隔。这里的token_type_ids会自动生成第一段全为 0第二段和第三段全为 1。这种设计让模型能区分“正文”和“实体名”比直接把实体拼进句子效果更稳定。only_first的含义是只对第一段即text做截断因为实体名通常很短不能截断。如果实体名过长导致超出 max_length可以改成truncationonly_second但 DUIE 实体一般不会那么长。2.3 loader 返回结构对训练循环的影响loader.py通常要返回一个 PyTorchDataset配合DataLoader使用。数据量不大时直接在__getitem__里做 tokenizer 编码也够用数据上十万条后最好在初始化时预先编码成input_ids数组避免训练时反复调 tokenizer 拖慢速度。我在实际项目中会把 loader 设计成支持两种模式训练模式做随机负采样验证模式固定负样本集。因为验证集如果每次负采样都变指标就无法复现。这个项目里没有分 dev 和 test但你自己扩展时一定要按住这个原则验证集必须静态。3. 模型构建BERT 编码加分类头以及为什么不用序列标注3.1 模型结构选择项目的 model.py 基于BertForSequenceClassification或自建BertModel Linear。基于BertForSequenceClassification最省事它内部已经在[CLS]输出上接了 dropout 和线性层。我们的关系分类任务有三个特性关系数量有限DUIE 常见 50 个以内、输入包含两个实体、判定是语义关系而非物理位置。因此分类头用一层就够没必要上 MLP 多层。如果自建模型典型结构如下class BertRelationModel(nn.Module): def __init__(self, num_labels, model_pathbert-base-chinese): super().__init__() self.bert BertModel.from_pretrained(model_path) self.dropout nn.Dropout(0.1) self.classifier nn.Linear(self.bert.config.hidden_size, num_labels) def forward(self, input_ids, token_type_ids, attention_mask): outputs self.bert( input_idsinput_ids, token_type_idstoken_type_ids, attention_maskattention_mask ) pooled outputs.pooler_output # 对应 [CLS] 的池化结果 logits self.classifier(self.dropout(pooled)) return logitspooler_output是 BERT 中[CLS]经过全连接和 Tanh 激活后的向量。它保留了整句语义也包含了实体对的信息因为实体文本已经拼进了输入序列。如果你需要显式标注实体的范围更精细的做法是取两个实体对应位置的平均 hidden state 再拼接或者加实体 marker token但这超出了本项目的范围。3.2 和序列标注 / pipeline 式抽取的差异早期实体关系抽取常用 pipeline先用 NER 模型抽实体再对实体对做分类。这种流程的坏处是误差传播——NER 漏抽实体关系模块就无从判断。BERT 关系分类模型虽然也是先给实体但实体可以由规则或别的方式获得也可以把 NER 和关系分类搭成两条并行线最后再匹配。DUIE 任务本身要求从一个句子中抽全体、主语、宾语和关系严格来说是联合抽取。但这份代码只负责“已知实体对判断关系”这一环属于 pipeline 中的第二步。别混淆这两个范围。在训练时标签类别数不要搞错。如果关系类型有 48 种再加上“无关系”就是 49 类。num_labels必须等于 49。常见错误是直接数spo_list里出现的 predicate 数量忘掉负样本类别导致维度不匹配。4. 训练验证与踩坑参数、脚本和损失函数细节4.1 损失函数与评估指标关系分类是标准的多分类问题用交叉熵损失即可。PyTorch 里nn.CrossEntropyLoss自带 softmax不需要在模型输出里再做。如果正负样本极不平衡可以给 loss 传weight参数让“无关系”类权重降低。DUIE 的负采样比例如果控制在 1:1 到 1:3不调权重也能收敛。评估指标我一般看 macro-F1因为每个关系的样本数不均micro-F1 容易被高频关系带跑。下面是常用的评估函数在 test.py 或验证脚本里可以用到from sklearn.metrics import classification_report def evaluate(model, dataloader, id2rel): model.eval() preds, trues [], [] with torch.no_grad(): for batch in dataloader: logits model(**batch) pred logits.argmax(dim-1).cpu().tolist() preds.extend(pred) trues.extend(batch[label].tolist()) print(classification_report(trues, preds, target_names[id2rel[i] for i in range(len(id2rel))]))注意classification_report输出的support列能反映每类样本量当某个关系 support 低于 10 时F1 没有意义你需要考虑数据增强或合并关系。4.2 训练超参数与 run.sh 脚本项目要求的环境是pytorch1.5.1 transformers2.5.1。这个版本组合下BertForSequenceClassification的 API 与现在transformers 4.x略有差异from_pretrained会自动下载模型但网络受限时要把bert-base-chinese目录放在本地并让from_pretrained指向那个目录。项目里 run.sh 大概长这样#!/bin/bash export CUDA_VISIBLE_DEVICES0 python main.py \ --model_path bert-base-chinese \ --train_path train.json \ --batch_size 16 \ --lr 2e-5 \ --epochs 3 \ --max_seq_len 128 \ --num_labels 50 \ --output_dir ./output我把num_labels示例写成了 50你务必根据自己数据集里的关系数加一。学习率 2e-5 是 BERT fine-tune 的常规起点。batch size 在显存够的情况下尽量用 16 或 32太小会导致 BN 类的层不稳定BERT 没有 BN 但分类头有 dropoutbatch 太小时梯度噪声大。epochs 我一般先跑 3观察 dev loss 是否回升。如果训练集只有几千条3 轮可能过拟合要提前早停。4.3 常见训练坑与排查路径下面表格是我跑完这个项目后归纳的高频错误现象原因排查与解法loss 不降学习率太大或标签类别数不匹配检查num_labels改成1e-5再试显存溢出max_seq_len 太长、batch_size 太大把序列长度减到 96batch 降到 8tokenizer 报错transformers 版本与代码不兼容按项目要求装transformers2.5.1验证集 F1 为 0负样本全部预测成了无关系或 id 映射错位打印id2rel逐条预测核对训练极慢每步都做 tokenizer 编码在__init__里预先 encode 所有样本另外nohup.out和test.sh出现在项目文件中说明原作者是通过 nohup 后台跑训练的。如果你在本地复现建议把print换成logging并周期性保存 checkpoint。以epoch为单位保存即可torch.save(model.state_dict(), foutput/bert_rel_{epoch}.pt)加载时要注意是state_dict还是完整模型。项目里 demo.py 通常用model.load_state_dict(torch.load(...))加载。如果保存的是整个模型直接torch.load会同时加载类定义容易因环境变更报错。5. 推理 demo 与工程化技巧把模型用到真实文本上5.1 给一段文本和两个实体输出关系demo.py的核心逻辑与训练时构造样本完全一致。唯一区别是推理时不需要标签并且要对同一对实体可能存在的多个token_type_ids保持确定性。下面是可运行的推理函数def predict(text, subject, object_, model, tokenizer, id2rel): encoded tokenizer.encode_plus( text, subject, object_, add_special_tokensTrue, max_length128, paddingmax_length, truncationonly_first, return_tensorspt ) logits model( input_idsencoded[input_ids].cuda(), token_type_idsencoded[token_type_ids].cuda(), attention_maskencoded[attention_mask].cuda() ) pred_id logits.argmax(dim-1).item() return id2rel[pred_id] # 示例 print(predict(《红楼梦》的作者是曹雪芹, 红楼梦, 曹雪芹, model, tokenizer, id2rel))注意这里subject和object_必须能在text中找到且顺序不能反。顺序反了模型输出的关系含义会变成“曹雪芹 与 红楼梦 的关系”预测结果大概率变成“无关系”。5.2 工程化的两个技巧缓存编码结果和服务化如果你要把模型部署成 HTTP 服务不要再重复调用generate之类全流程。我的做法是把 tokenizer 和模型加载一次之后每次请求只做一次encode_plus和一次 forward。这个项目的 demo.py 是单条调用足够展示功能。但如果文本量大可以把句子预先 tokenize 好并按 batch 推理。记住一点推理时padding建议设为max_length而不是longest之外的动态 padding否则 batch 内序列长度不一致GPU 利用率会下降。另一个技巧是把id2rel映射存成 JSON和模型权重一起打包。千万不要在部署环境里从训练集重新生成映射否则一旦训练集顺序变化同一 id 对应的关系就变了线上预测会静默出错。5.3 关系抽取结果如何接入知识图谱模型输出的是关系名称但要写入图谱还需要把实体和关系组成三元组。常见做法是在拿到predicate后校验该关系是否在 schema 允许的范围内再写入 Neo4j 或图数据库。如果多个关系同时出现比如“作者”和“朝代”你需要对同一对实体做多标签预测而不是简单 argmax。实际处理时可以把分类头从nn.Linear换成nn.Linear Sigmoid损失函数换成BCEWithLogitsLoss。项目给出的代码只支持单标签但工程场景里多标签更常见升级方式就是改最后两层。李沐老师在 BERT 精读课里强调过预训练模型的价值在于对上下文的建模能力。用在实体关系抽取上你不需要手工设计实体间的距离特征或句法路径只要把实体文本拼进输入模型就能从词语共现和历史知识中捕捉到关系。所谓踩坑其实大多数集中在数据构造的一致性上训练时怎么截断推理时也要同样截断训练时怎么设定token_type_ids推理时也要用同一套模板。做到这一点整个流程就跑得通。本文还有配套的精品资源点击获取 
长河君 · 2026/9/17 1:51:08 · 阅读 1.2k 基于BERT的实体关系抽取实战：DUIE数据集与关系分类
简介基于PyTorch与BERT的实体关系抽取完整工程代码包面向自然语言处理初学者与知识图谱构建者适用于中文关系抽取任务。项目中已内置DUIE数据集加载逻辑支持通过配置文件替换为自定义数据可快速跑通训练、测试与推理流程基于PyTorch 1.5.1与Transformers 2.5.1环境测试通过。压缩包共20个文件包括5个Python源码、5个编译缓存、5张结果展示图、2个Shell启动脚本及运行日志整体仅2.45MB轻量易用Python脚本覆盖数据加载、模型定义、训练、测试与演示等关键环节Shell脚本可一键启动训练和测试另附可视化图片便于直观查看效果整体目录结构清晰、模块职责分明便于二次开发已有248人学习下载。阅读源码后可掌握BERT序列标注思路、关系分类训练细节、自定义数据接入方法以及Shell脚本执行流程适合作为入门知识图谱关系抽取的参考项目利用该方案可快速复现关系抽取实验为后续知识图谱构建打下基础。1. 实体关系抽取为什么要落在 BERT 上实体关系抽取是知识图谱构建里最费人力的环节。早期做法靠特征模板写正则匹配“位于”“出生于”这类触发词换领域就失效后来用 BiLSTM-CRF 做序列标注但关系类型一多标签体系爆炸而且无法处理跨长距离的实体依赖。用 BERT 做这件事的思路不同把关系分类变成一句“句子 头实体 尾实体”的语义匹配问题让预训练模型去判断中间有没有关系、是什么关系。DUIE 是百度信息抽取大赛的数据集里面有 schema、spo_list正好拿来做关系抽取的训练与评测。这个项目的主线是用 DUIE 数据构造关系分类样本用 bert-base-chinese 做编码接一个线性分类头输出关系类型。对 5 年以上的工程师来说值得关注的是数据对齐和采样策略而不是模型本身因为 BERT 部分已经很成熟坑大多在教学流程之外。2. 数据准备与 loader 设计从 DUIE 原始结构到训练样本2.1 DUIE 的原始格式与本次任务的差异DUIE 官方数据是 JSON 行每条包含text和spo_list。spo_list里每一项是predicate、subject、object以及它们在文本中的偏移量。下面的片段是原始格式的典型样子{ text: 《红楼梦》的作者是曹雪芹曹雪芹名霑字梦阮。, spo_list: [ { predicate: 作者, subject: 红楼梦, object: 曹雪芹, subject_offset: [0, 4], object_offset: [9, 12] } ] }注意本项目里的train.json就是这种结构。但我们要做的是关系分类不是序列标注所以需要从spo_list中抽出正样本再把不含该关系的实体对构造为负样本。关系分类的输入通常是句子 头实体 尾实体输出是关系类型。如果直接沿用原始 JSON模型会遇到两个问题一条文本里可能有多对实体、多个关系直接整句分类无法确定“谁和谁”的关系。因此 loader 的核心职责是样本构造和 token 对齐。2.2 构造关系分类样本的策略常见做法是遍历spo_list中的每个三元组把subject当成头实体object当成尾实体predicate当作标签与此同时随机抽取文本中其他实体组合作为负样本。这个项目里 loader.py 也是这么处理的具体分三步2.2.1 读取与字段解析先按行加载 JSON解析出text与spo_list。对每条 SPO生成一条正样本(text, subject, object, predicate_id)。如果某条文本同时有多个 SPO就生成多条样本保证每个三元组都能被训练到。2.2.2 负样本采样正样本的数量通常远少于错配的实体对数量。为了模型学会区分“有关系”与“无关系”需要负样本。我一般用的采样方式是固定一条文本随机选择一个出现在该文本中的实体作为头再随机选择另一个实体作为尾如果这个组合不在spo_list中就标记为“无关系”类别比如predicate为NEG。如果采样后发现正负比例超过 1:3就减少负采样数量避免模型偏向多数类。2.2.3 文本拼接与 tokenizer 编码BERT 不能直接吃原始三元组需要按固定模板拼成一个序列。项目里用的是“[CLS] 文本 [SEP] 头实体 [SEP] 尾实体 [SEP]”的格式。用tokenizer.encode_plus可以一次拿到input_ids、token_type_ids和attention_maskfrom transformers import BertTokenizer tokenizer BertTokenizer.from_pretrained(bert-base-chinese) def build_one_sample(text, subject, object_, label): encoded tokenizer.encode_plus( text, subject, object_, add_special_tokensTrue, max_length128, paddingmax_length, truncationonly_first, # 优先截断文本 return_tensorspt ) return { input_ids: encoded[input_ids].squeeze(0), token_type_ids: encoded[token_type_ids].squeeze(0), attention_mask: encoded[attention_mask].squeeze(0), label: label }注意encode_plus传三个参数时BERT 会把第二、第三个参数当作text_pair并用两个或多个[SEP]分隔。这里的token_type_ids会自动生成第一段全为 0第二段和第三段全为 1。这种设计让模型能区分“正文”和“实体名”比直接把实体拼进句子效果更稳定。only_first的含义是只对第一段即text做截断因为实体名通常很短不能截断。如果实体名过长导致超出 max_length可以改成truncationonly_second但 DUIE 实体一般不会那么长。2.3 loader 返回结构对训练循环的影响loader.py通常要返回一个 PyTorchDataset配合DataLoader使用。数据量不大时直接在__getitem__里做 tokenizer 编码也够用数据上十万条后最好在初始化时预先编码成input_ids数组避免训练时反复调 tokenizer 拖慢速度。我在实际项目中会把 loader 设计成支持两种模式训练模式做随机负采样验证模式固定负样本集。因为验证集如果每次负采样都变指标就无法复现。这个项目里没有分 dev 和 test但你自己扩展时一定要按住这个原则验证集必须静态。3. 模型构建BERT 编码加分类头以及为什么不用序列标注3.1 模型结构选择项目的 model.py 基于BertForSequenceClassification或自建BertModel Linear。基于BertForSequenceClassification最省事它内部已经在[CLS]输出上接了 dropout 和线性层。我们的关系分类任务有三个特性关系数量有限DUIE 常见 50 个以内、输入包含两个实体、判定是语义关系而非物理位置。因此分类头用一层就够没必要上 MLP 多层。如果自建模型典型结构如下class BertRelationModel(nn.Module): def __init__(self, num_labels, model_pathbert-base-chinese): super().__init__() self.bert BertModel.from_pretrained(model_path) self.dropout nn.Dropout(0.1) self.classifier nn.Linear(self.bert.config.hidden_size, num_labels) def forward(self, input_ids, token_type_ids, attention_mask): outputs self.bert( input_idsinput_ids, token_type_idstoken_type_ids, attention_maskattention_mask ) pooled outputs.pooler_output # 对应 [CLS] 的池化结果 logits self.classifier(self.dropout(pooled)) return logitspooler_output是 BERT 中[CLS]经过全连接和 Tanh 激活后的向量。它保留了整句语义也包含了实体对的信息因为实体文本已经拼进了输入序列。如果你需要显式标注实体的范围更精细的做法是取两个实体对应位置的平均 hidden state 再拼接或者加实体 marker token但这超出了本项目的范围。3.2 和序列标注 / pipeline 式抽取的差异早期实体关系抽取常用 pipeline先用 NER 模型抽实体再对实体对做分类。这种流程的坏处是误差传播——NER 漏抽实体关系模块就无从判断。BERT 关系分类模型虽然也是先给实体但实体可以由规则或别的方式获得也可以把 NER 和关系分类搭成两条并行线最后再匹配。DUIE 任务本身要求从一个句子中抽全体、主语、宾语和关系严格来说是联合抽取。但这份代码只负责“已知实体对判断关系”这一环属于 pipeline 中的第二步。别混淆这两个范围。在训练时标签类别数不要搞错。如果关系类型有 48 种再加上“无关系”就是 49 类。num_labels必须等于 49。常见错误是直接数spo_list里出现的 predicate 数量忘掉负样本类别导致维度不匹配。4. 训练验证与踩坑参数、脚本和损失函数细节4.1 损失函数与评估指标关系分类是标准的多分类问题用交叉熵损失即可。PyTorch 里nn.CrossEntropyLoss自带 softmax不需要在模型输出里再做。如果正负样本极不平衡可以给 loss 传weight参数让“无关系”类权重降低。DUIE 的负采样比例如果控制在 1:1 到 1:3不调权重也能收敛。评估指标我一般看 macro-F1因为每个关系的样本数不均micro-F1 容易被高频关系带跑。下面是常用的评估函数在 test.py 或验证脚本里可以用到from sklearn.metrics import classification_report def evaluate(model, dataloader, id2rel): model.eval() preds, trues [], [] with torch.no_grad(): for batch in dataloader: logits model(**batch) pred logits.argmax(dim-1).cpu().tolist() preds.extend(pred) trues.extend(batch[label].tolist()) print(classification_report(trues, preds, target_names[id2rel[i] for i in range(len(id2rel))]))注意classification_report输出的support列能反映每类样本量当某个关系 support 低于 10 时F1 没有意义你需要考虑数据增强或合并关系。4.2 训练超参数与 run.sh 脚本项目要求的环境是pytorch1.5.1 transformers2.5.1。这个版本组合下BertForSequenceClassification的 API 与现在transformers 4.x略有差异from_pretrained会自动下载模型但网络受限时要把bert-base-chinese目录放在本地并让from_pretrained指向那个目录。项目里 run.sh 大概长这样#!/bin/bash export CUDA_VISIBLE_DEVICES0 python main.py \ --model_path bert-base-chinese \ --train_path train.json \ --batch_size 16 \ --lr 2e-5 \ --epochs 3 \ --max_seq_len 128 \ --num_labels 50 \ --output_dir ./output我把num_labels示例写成了 50你务必根据自己数据集里的关系数加一。学习率 2e-5 是 BERT fine-tune 的常规起点。batch size 在显存够的情况下尽量用 16 或 32太小会导致 BN 类的层不稳定BERT 没有 BN 但分类头有 dropoutbatch 太小时梯度噪声大。epochs 我一般先跑 3观察 dev loss 是否回升。如果训练集只有几千条3 轮可能过拟合要提前早停。4.3 常见训练坑与排查路径下面表格是我跑完这个项目后归纳的高频错误现象原因排查与解法loss 不降学习率太大或标签类别数不匹配检查num_labels改成1e-5再试显存溢出max_seq_len 太长、batch_size 太大把序列长度减到 96batch 降到 8tokenizer 报错transformers 版本与代码不兼容按项目要求装transformers2.5.1验证集 F1 为 0负样本全部预测成了无关系或 id 映射错位打印id2rel逐条预测核对训练极慢每步都做 tokenizer 编码在__init__里预先 encode 所有样本另外nohup.out和test.sh出现在项目文件中说明原作者是通过 nohup 后台跑训练的。如果你在本地复现建议把print换成logging并周期性保存 checkpoint。以epoch为单位保存即可torch.save(model.state_dict(), foutput/bert_rel_{epoch}.pt)加载时要注意是state_dict还是完整模型。项目里 demo.py 通常用model.load_state_dict(torch.load(...))加载。如果保存的是整个模型直接torch.load会同时加载类定义容易因环境变更报错。5. 推理 demo 与工程化技巧把模型用到真实文本上5.1 给一段文本和两个实体输出关系demo.py的核心逻辑与训练时构造样本完全一致。唯一区别是推理时不需要标签并且要对同一对实体可能存在的多个token_type_ids保持确定性。下面是可运行的推理函数def predict(text, subject, object_, model, tokenizer, id2rel): encoded tokenizer.encode_plus( text, subject, object_, add_special_tokensTrue, max_length128, paddingmax_length, truncationonly_first, return_tensorspt ) logits model( input_idsencoded[input_ids].cuda(), token_type_idsencoded[token_type_ids].cuda(), attention_maskencoded[attention_mask].cuda() ) pred_id logits.argmax(dim-1).item() return id2rel[pred_id] # 示例 print(predict(《红楼梦》的作者是曹雪芹, 红楼梦, 曹雪芹, model, tokenizer, id2rel))注意这里subject和object_必须能在text中找到且顺序不能反。顺序反了模型输出的关系含义会变成“曹雪芹 与 红楼梦 的关系”预测结果大概率变成“无关系”。5.2 工程化的两个技巧缓存编码结果和服务化如果你要把模型部署成 HTTP 服务不要再重复调用generate之类全流程。我的做法是把 tokenizer 和模型加载一次之后每次请求只做一次encode_plus和一次 forward。这个项目的 demo.py 是单条调用足够展示功能。但如果文本量大可以把句子预先 tokenize 好并按 batch 推理。记住一点推理时padding建议设为max_length而不是longest之外的动态 padding否则 batch 内序列长度不一致GPU 利用率会下降。另一个技巧是把id2rel映射存成 JSON和模型权重一起打包。千万不要在部署环境里从训练集重新生成映射否则一旦训练集顺序变化同一 id 对应的关系就变了线上预测会静默出错。5.3 关系抽取结果如何接入知识图谱模型输出的是关系名称但要写入图谱还需要把实体和关系组成三元组。常见做法是在拿到predicate后校验该关系是否在 schema 允许的范围内再写入 Neo4j 或图数据库。如果多个关系同时出现比如“作者”和“朝代”你需要对同一对实体做多标签预测而不是简单 argmax。实际处理时可以把分类头从nn.Linear换成nn.Linear Sigmoid损失函数换成BCEWithLogitsLoss。项目给出的代码只支持单标签但工程场景里多标签更常见升级方式就是改最后两层。李沐老师在 BERT 精读课里强调过预训练模型的价值在于对上下文的建模能力。用在实体关系抽取上你不需要手工设计实体间的距离特征或句法路径只要把实体文本拼进输入模型就能从词语共现和历史知识中捕捉到关系。所谓踩坑其实大多数集中在数据构造的一致性上训练时怎么截断推理时也要同样截断训练时怎么设定token_type_ids推理时也要用同一套模板。做到这一点整个流程就跑得通。本文还有配套的精品资源点击获取