简介:面向自然语言处理与深度学习研究者的中文医学实体关系抽取实现,基于CMeIE-V2和CMeIE两个数据集完成数据清洗、标注与评测,统一实现了卷积神经网络(CNN)、循环神经网络(RNN)、胶囊网络(Capsule)、基于剪枝依赖树的图卷积网络(GCN)、Transformer与BERT六类模型的训练和推理流程。与常规文本分类不同,项目在关系分类中引入上下文文本、实体信息、实体位置编码及实体类别特征,可帮助读者理解如何融合多种结构化信号提升关系抽取效果。压缩包共484个文件、约21.7MB,以249个Python源码文件为主,辅以182个pyc字节码、25个YAML配置、7个CSV数据文件及6个Markdown说明文档,便于直接阅读、调试与复现实验。已有200人学习,适合希望对比经典模型与预训练模型在中文医学关系抽取任务上表现的开发者、研究生及竞赛团队,可直接基于现有代码扩展新的数据集或模型结构。
1. 关系抽取里的六类模型同台竞技:这套代码在解决什么
关系抽取是信息抽取里最费人工的一环,给定一个句子和两个实体,要自动判断实体对属于哪种关系。市面上讲 CNN、RNN、Transformer 的教程很多,但能把 CNN、RNN、Capsule、GCN、Transformer 和 BERT 六类模型放进同一套框架里对比的代码包很少见。这份资源的价值在于数据预处理、实体遮蔽、模型统一接口、训练评估全部打通,同一份标准数据集上六种架构直接出 F1。适合两类人:想尽快跑通关系抽取完整流程的新手,以及要做模型选型消融对比的工程师。我拆完这套代码后,把里面的预处理细节和排错经验整理成下面的笔记。
2. 数据预处理:实体遮蔽、相对位置与句法邻接矩阵
关系抽取的数据处理和文本分类不太一样。文本分类只需要「句子→标签」,关系抽取必须带实体位置信息,否则模型根本不知道要判断哪两个词之间的关系。这一步做不干净,后面六个模型再花哨也白搭。
2.1 统一样本结构:先定数据格式再谈模型
这套代码的原始数据走的是 TSV 格式,一行一个样本,六列:分词后的句子、第一个实体的起始位置和结束位置、第二个实体的起始位置和结束位置、关系标签。没有关系的样本标签记为NA,这种负样本在关系抽取里很重要,后面评估时能明显拉开模型差距。
代码开头先定义一个统一的样本结构:
@dataclass class Sample: tokens: List[str] # 句子 token 列表 e1_pos: List[int] # 第一个实体的 token 起始/结束下标 e2_pos: List[int] # 第二个实体的 token 起始/结束下标 label: str # 关系类型,NA 表示无关系 def parse_line(line: str) -> Sample: parts = line.strip().split("\t") # 输入格式: tokens \t e1_start \t e1_end \t e2_start \t e2_end \t label return Sample( tokens=parts[0].split(), e1_pos=[int(parts[1]), int(parts[2])], e2_pos=[int(parts[3]), int(parts[4])], label=parts[5], )这里把实体位置直接存成 token 下标,而不是字符偏移,原因在于后面的遮蔽和位置编码都基于 token 粒度操作。如果你的原始数据给的是字符级位置,通常的做法是先用分词器把句子切好,再做一次字符偏移到 token 下标的映射。这套代码假定输入已经是 token 级,省了这一步。
2.2 实体遮蔽与相对位置编码
不管用 CNN 还是 BERT,模型都需要知道两个实体的边界在哪。最常见的做法是实体遮蔽(mark entity):把第一个实体的全部 token 替换成[E1],第二个实体的全部 token 替换成[E2]。这样长实体也不需要模型去额外推理边界。
def mark_entities(sample: Sample): """实体位置替换成占位符,保留实体类型信息""" tokens = sample.tokens[:] e1_s, e1_e = sample.e1_pos e2_s, e2_e = sample.e2_pos tokens[e1_s:e1_e] = ["[E1]"] tokens[e2_s:e2_e] = ["[E2]"] return tokens遮蔽之后还需要相对位置特征。Transformer 系模型靠位置编码,但位置编码是绝对位置,对实体间距离不敏感。我一般会额外拼两个位置序列:每个 token 到第一个实体的距离、到第二个实体的距离,截断到[-max_len, max_len],再加一个偏置转为非负下标,查 embedding 表拿到位置向量。
def pos_features(sample: Sample, max_len=100): e1_s = sample.e1_pos[0] e2_s = sample.e2_pos[0] pos1 = [i - e1_s for i in range(len(sample.tokens))] pos2 = [i - e2_s for i in range(len(sample.tokens))] # 截断到 [-max_len, max_len],加 max_len 保证下标非负 pos1 = [min(max(p, -max_len), max_len) + max_len for p in pos1] pos2 = [min(max(p, -max_len), max_len) + max_len for p in pos2] return pos1, pos2这个截断参数max_len=100是个经验值。绝大多数句子的实体距离都在 100 个 token 以内,超过这个范围的位置信息已经几乎没有区分度。如果你想压缩 embedding 表大小,可以把距离压缩到 60,效果下降不明显。
2.3 句法依赖邻接矩阵:GCN 模型的输入
GCN 和前面几个模型最大的区别是它需要图结构输入。关系抽取里这张图通常由句法依赖树构造:每个 token 是节点,依赖关系是边。核心操作是把边列表转成邻接矩阵,然后做归一化。这一步很容易出问题,我见过不少复现代码直接拿原始邻接矩阵乘特征,结果深层的梯度全炸了。
def dependency_adjacency(edges, num_tokens): """依赖边列表转邻接矩阵,加自环并做对称归一化""" adj = np.zeros((num_tokens, num_tokens)) for head, dep in edges: # 依赖树是有向的,但 GCN 聚合需要无向传播 adj[head, dep] = 1.0 adj[dep, head] = 1.0 for i in range(num_tokens): adj[i, i] = 1.0 # 自环,保留节点自身信息 # 对称归一化: D^-1/2 A D^-1/2 degree = adj.sum(axis=1) degree_inv_sqrt = np.power(degree + 1e-6, -0.5) degree_inv_sqrt = np.diag(degree_inv_sqrt) adj_norm = degree_inv_sqrt @ adj @ degree_inv_sqrt return torch.tensor(adj_norm, dtype=torch.float32)注意边要转成无向图,因为 GCN 聚合邻居信息时,一个 token 既要接收依赖它的子节点信息,也要接收它依赖的父节点信息。句法解析这一步通常用现成的依赖解析工具生成(head, dep)边列表,如果你拿到的数据没有句法标注,需要额外跑一遍解析,这会明显增加预处理时间,我一般会缓存解析结果到本地文件,避免每次训练重复解析。
2.4 BERT 的子词对齐:最隐蔽的坑
BERT 类预训练模型用的是 WordPiece 分词,一个 token 会被切碎成多个子词,实体位置和句子长度全部对不上。直接按原始 token 下标取嵌入,拿到的就是错位信息。代码里要维护一个word_ids映射,把每个子词映射回原始 token 下标,实体边界按子词重新展开。
def bertize(sample: Sample, tokenizer): # 按 token 序列做分词,is_split_into_words 防止整句被再切一次 encoding = tokenizer( sample.tokens, is_split_into_words=True, return_offsets_mapping=True, truncation=True, max_length=256, ) word_ids = encoding.word_ids() # 每个子词对应的原始 token 下标 e1_start = word_ids.index(sample.e1_pos[0]) # 实体最后一段子词的结束位置要从尾部找 e1_end = len(word_ids) - 1 - word_ids[::-1].index(sample.e1_pos[1] - 1) return encoding, e1_start, e1_end这里有个细节:找结束位置要从word_ids尾部反向找。因为实体的最后一个 token 可能被切出多个子词,正向index永远会命中第一个子词,拿不到真正的尾部边界。这个反向查找我用过一次后就成了固定写法。
3. 六类模型实现骨架:从 TextCNN 到 BERT 的统一接口
六类模型结构差异很大,但训练循环、评估逻辑完全一致。这套代码把模型全部包成同一个BaseModel接口,前向只返回 logits,loss 在内部算,外层训练代码不用管每个模型内部的细节。这个抽象选得很实用,替换模型时只改一行注册代码。
3.1 统一接口与训练循环约定
class BaseModel(nn.Module): """所有模型的统一接口,forward 返回 logits,loss 在内部计算""" def __init__(self, num_labels, **kwargs): super().__init__() self.num_labels = num_labels def forward(self, batch): # 子类必须实现 raise NotImplementedError def compute_loss(self, logits, labels): return F.cross_entropy(logits, labels, ignore_index=-100)训练循环里只看compute_loss的返回值,六种模型的 loss 计算方式都一样是交叉熵。ignore_index=-100是预训练模型里的常见约定,用来屏蔽 padding 位置;如果你的数据里没有 padding token,这个参数不生效,但保留无害。我习惯在注册表里集中管理六个模型的构建入口,新增模型时只需要往字典里加一个类。
3.2 TextCNN:用多尺度卷积抓局部 n-gram
TextCNN 在关系抽取里不算最强,但作为 baseline 足够快。核心思想是用多个不同宽度的卷积核并行抽取 n-gram 特征。宽度 3、4、5 分别对应三元组、四元组、五元组特征。
class TextCNN(BaseModel): def __init__(self, embedding_dim=300, num_filters=128, kernel_sizes=[3, 4, 5], num_labels=19): super().__init__(num_labels) self.convs = nn.ModuleList([ nn.Conv1d(embedding_dim, num_filters, k, padding=k // 2) for k in kernel_sizes ]) self.classifier = nn.Linear(len(kernel_sizes) * num_filters, num_labels) def forward(self, batch): # x: [batch_size, seq_len, embedding_dim] x = batch["embedding"] x = x.transpose(1, 2) # Conv1d 需要 [batch, channels, seq_len] pooled = [] for conv in self.convs: h = torch.relu(conv(x)) pooled.append(torch.max(h, dim=2).values) out = torch.cat(pooled, dim=1) return self.classifier(out)kernel_sizes=[3,4,5]是个很稳的起点,实体对之间的局部线索大多落在三到五个 token 的窗口内。如果你处理的是长尾句式,可以加一个宽度为 7 的卷积核,但参数量会涨,在小数据集上容易过拟合。
3.3 BiLSTM:序列编码加实体头向量
BiLSTM 的做法是把整个句子编码成隐状态序列,然后取两个实体位置上的隐状态拼接,或者像注意力池化那样加权汇总。这套代码用的是取实体头 token 隐状态的方式,简单而且效果不差。
class BiLSTM(BaseModel): def __init__(self, embedding_dim=300, hidden_dim=256, num_layers=2, dropout=0.3, num_labels=19): super().__init__(num_labels) self.lstm = nn.LSTM( embedding_dim, hidden_dim, num_layers=num_layers, bidirectional=True, batch_first=True, dropout=dropout, ) self.classifier = nn.Linear(hidden_dim * 4, num_labels) def forward(self, batch): x = batch["embedding"] e1_head = batch["e1_head"] # 实体起始位置 e2_head = batch["e2_head"] h, _ = self.lstm(x) h_e1 = h[torch.arange(h.size(0)), e1_head] h_e2 = h[torch.arange(h.size(0)), e2_head] out = torch.cat([h_e1, h_e2], dim=1) return self.classifier(out)注意hidden_dim * 4的来历:双向 LSTM 每个方向输出hidden_dim,两个方向拼接是2 * hidden_dim,两个实体各取一份再拼接就是 4 倍。这个数字经常有人写错,少算了一倍维度,训练时直接报 shape 不匹配。num_layers=2足够,加深到 3 层以上在中小规模数据上收益很小。
3.4 Capsule:动态路由里的 squash 函数
Capsule 网络在关系抽取里的角色比较特殊,它的动态路由机制理论上能捕捉实体间的组合特征。这套代码的实现是标准的 capsule 卷积层加squash激活:
def squash(x, dim=-1): """胶囊网络的压缩函数,把向量模长压到 [0, 1)""" norm = torch.norm(x, dim=dim, keepdim=True) scale = norm / (1 + norm ** 2) return scale * x def routing(b_ij, u_hat, num_iter=3): """动态路由:迭代更新耦合系数 c_ij""" for _ in range(num_iter): c_ij = torch.softmax(b_ij, dim=1) s_j = (c_ij * u_hat).sum(dim=1) v_j = squash(s_j) b_ij = b_ij + (u_hat * v_j.unsqueeze(1)).sum(dim=-1) return v_jnum_iter=3是原论文的经验值。路由迭代次数不是越多越好,我实测过迭代 5 次以上训练会变慢,F1 基本不涨。squash里的1 + norm**2变体比较常见,也有人用norm + 1e-6的版本,效果差别不大,但梯度稳定性更好,你可以两个都试。
3.5 GCN:句法图上做邻居聚合
GCN 层负责在句法依赖树上传播信息,用的就是第 2 节生成的归一化邻接矩阵。实现上比较直接:
class GCNLayer(nn.Module): def __init__(self, in_dim, out_dim, dropout=0.2): super().__init__() self.weight = nn.Parameter(torch.randn(in_dim, out_dim)) self.dropout = nn.Dropout(dropout) def forward(self, x, adj): # x: [batch, seq_len, in_dim] # adj: [batch, seq_len, seq_len] support = torch.matmul(x, self.weight) out = torch.matmul(adj, support) return torch.relu(self.dropout(out))GCN 层数一般是 2 到 3 层。层数太多会把句法树的全局信息过度平滑,每个节点的表示趋同,反而丢掉了实体间的局部差异。做关系抽取时我通常只在句法依赖图上叠两层 GCN,第三层开始 F1 回落。
3.6 Transformer 与 BERT:自注意力与预训练微调
Transformer 实现用的是标准编码器:
encoder_layer = nn.TransformerEncoderLayer( d_model=256, nhead=8, dim_feedforward=1024, dropout=0.1, batch_first=True, ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=3)nhead=8要求d_model能整除 8,改成 6 或者 4 时记得同步调整d_model,否则初始化报错。BERT 微调则是加载预训练模型,把[E1][E2]两个特殊标记扩进 embedding 表,只对新增的两个向量做随机初始化:
from transformers import AutoConfig, AutoModel def build_bert_model(model_path, num_labels): config = AutoConfig.from_pretrained(model_path, num_labels=num_labels) model = AutoModel.from_pretrained(model_path, config=config) # 扩展 embedding 表,加入 [E1] [E2] model.resize_token_embeddings(model.config.vocab_size + 2) return modelresize_token_embeddings会自动填充随机向量。这里有一个细节:新增 token 的 embedding 随机初始化会导致最初几个 step 的 loss 异常高,因为这两个向量没有预训练语义。正常现象,一般跑几百个 step 后就会被拉回来。
六类模型的核心对比如下:
| 模型 | 核心机制 | 关键超参数 | 适合处理的信息 |
|---|---|---|---|
| TextCNN | 多尺度卷积 n-gram | kernel_sizes=[3,4,5] | 局部词序模式 |
| BiLSTM | 双向序列编码 | hidden_dim=256, num_layers=2 | 中长距离序列 |
| Capsule | 动态路由组合特征 | num_iter=3 | 组合语义特征 |
| GCN | 句法图邻居聚合 | num_layers=2 | 依赖树结构信息 |
| Transformer | 自注意力全连接 | num_layers=3, nhead=8 | 全局依赖 |
| BERT | 预训练微调 | 最大长度、学习率 | 深度语义表示 |
4. 避坑与常见问题:五组训练翻车现场
下面这几组问题是我在这类关系抽取项目里反复踩过的,每条都按现象、原因、解决三个步骤拆开,方便你对照排查。
4.1 BERT 加了特殊标记后 loss 不下降
现象:把[E1][E2]加入 tokenizer 后,训练两三个 epoch,验证 F1 一直停在很低的水平,loss 也降得很慢。
原因:[E1]和[E2]不在预训练词表里,resize_token_embeddings给的是随机初始化向量,而这个向量要同时承担「实体边界标记」和「语义内容」两种职责。更隐蔽的是,如果你没有把实体类型信息并入 token embedding,模型根本不知道这是实体标记还是普通单词。
解决:两件事一起做。第一,用resize_token_embeddings扩展词表后,单独放大新增 embedding 的初始化尺度,比如用标准差 0.1 的正态分布,而不是默认的 0.02。第二,把实体类型信息拼接到 BERT 输出之上,比如[E1]位置对应的隐状态再接一个实体类型向量,再送进分类层。做完整合后 F1 基本能回到预期水平。
4.2 GCN 深层训练时梯度爆炸或消失
现象:GCN 堆到四层及以上,或者归一化实现不对时,训练 loss 直接变成 NaN,或者梯度范数暴涨。
原因:邻接矩阵没有做对称归一化。如果直接把原始 0/1 邻接矩阵和特征相乘,节点度数大的位置会累积大量数值,经过多层传播后尺度失控。另一个原因是 ReLU 负半轴把大量梯度截断,多层 GCN 里信息流衰减极快。
解决:必须使用第 2 节里的D^-1/2 A D^-1/2对称归一化,并且控制层数不超过 3 层。我排查 NaN 时习惯先打印每一层的输出范数,从第一层开始看哪一层开始发散;绝大多数情况是归一化忘掉了自环。
4.3 Capsule 的动态路由结果反复横跳
现象:同一个 batch,调一下随机种子,验证集 F1 能差 3 到 4 个点。路由迭代次数改了之后效果也很不稳定。
原因:动态路由里b_ij的初始值是随机的,迭代过程中它会根据胶囊之间的相似度被反复更新,这个过程对初始化敏感。我在代码里试过,num_iter从 2 调到 5,F1 波动超过 2 个点,这在对比实验里很难判断改进是真的还是噪声。
解决:固定全局随机种子,把num_iter=3定死在配置里不要随便动。如果你要对比不同模型的稳定性,每个配置至少要跑 3 个种子取平均值,只看单次结果很容易被带偏。Capsule 这部分本身就是六类模型里最「玄学」的一个,做好心理准备。
4.4 类别不平衡让 F1 虚高
现象:评估代码直接算 accuracy,测试集上所有NA样本都被分对,模型准确率到了 92%,看起来效果很好,但正类关系的 P/R/F1 一塌糊涂。
原因:公开关系抽取数据集的NA类别占比通常超过一半,模型只要无脑预测NA就能有很高的准确率。准确率在这个任务里完全失真,真正要看的是 macro-F1 或者按关系类别分开看的 F1。
解决:评估必须走分类报告那套逻辑。下面这段是我固定的评估函数:
from sklearn.metrics import classification_report def evaluate(model, loader): model.eval() preds, labels = [], [] with torch.no_grad(): for batch in loader: logits = model(batch) preds.extend(logits.argmax(dim=1).cpu().tolist()) labels.extend(batch["label"].cpu().tolist()) report = classification_report(labels, preds, digits=4, zero_division=0) return reportzero_division=0要加上,否则某个类别在预测中完全没出现时,classification_report会抛警告甚至报错。如果你发现正类 F1 明显偏低,优先考虑负采样,把NA样本砍到正类的两倍以内。
4.5 复现论文分数差一截
现象:拿着同样的超参数复现,跑出来的 F1 比论文报告的低 2 到 3 个点,怎么调都补不回来。
原因:关系抽取实验的方差比想象中大。随机种子、预训练模型的权重版本、最大序列长度、batch 大小、学习率调度策略,任何一个变量不同,最终 F1 都可能差出 2 个点。更隐蔽的是数据集的划分方式,很多公开数据集的官方划分和随机划分结果差异很大。
解决:我一般把实验配置全部固定成常量写在配置文件里,包括 seed、max_len、batch_size、学习率、warmup 步数。每次跑实验前先做一次 smoke test,用 200 条样本跑通流程,再上全量数据。报告分数时同时报三粒种子的均值±方差,不单挑最高的一次说事。
5. 把模型压到极限:小样本冒烟与超参消融的验证顺序
关系抽取模型的调参有不少玄学成分,但好的调试顺序能把变量控制住。我这套顺序是:先小样本冒烟,再超参消融,最后全量复现。
小样本冒烟的目标不是刷分,而是验证模型能不能过拟合。取 200 条训练样本,让模型跑 5 个 epoch,正常情况下训练 loss 应该下降到接近 0。如果 200 条样本都过拟合不了,说明模型构建有 bug,比如实体位置错位、标签对应错、embedding 没加载对。冒烟过了再谈调参。
冒烟代码固定几个关键参数:
def smoke_test(model, train_loader): """200 条样本跑 5 个 epoch,验证模型能过拟合""" optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3) for epoch in range(5): model.train() total_loss = 0.0 for batch in train_loader: logits = model(batch) loss = model.compute_loss(logits, batch["label"]) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"epoch {epoch}: loss={total_loss:.4f}")冒烟通过后做超参消融。优先消融的顺序是:实体遮蔽方式 → 位置特征维度 → 学习率 → GCN 层数 / Capsule 路由迭代数。学习率从2e-5到5e-4区间各跑一次,BERT 微调用小学习率,CNN/LSTM 这类从零训练的模型可以适当放大。
我给的参考基线如下:
| 参数 | 从零训练模型 | BERT 微调 |
|---|---|---|
| batch_size | 64 | 32 |
| 学习率 | 1e-3 | 2e-5 |
| 学习率调度 | 线性衰减 | warmup + 线性衰减 |
| 最大长度 | 128 | 256 |
| 训练轮数 | 15 | 5 |
从零训练的模型可以开大 batch size,因为梯度方差小,收敛更稳。BERT 类预训练模型 batch size 开大会把显存打满,而且微调阶段大 batch 容易让 embedding 层发生灾难性遗忘。
做消融时要盯住验证集 F1 的波动区间。每个配置跑三个种子,三个种子的最大值和最小值差距超过 1.5 个点,就说明这个超参数不是当前主要矛盾,优先去调别的变量。差距稳定在 0.5 个点以内,才判断改动方向有效。
全量复现阶段,我的习惯是把 data loader 的num_workers和预取(prefetch)开满,让 GPU 不空等。BERT 微调那一类模型,我还会加上梯度累积,模拟更大的 batch 而不牺牲显存。
最后说一个我自己的教训:以前拿到新的代码包,总喜欢直接全量训练看效果,结果十次里有六次浪费在跑 bug 上。从那以后,我每次换模型都强制走一遍「冒烟 → 超参消融 → 全量复现」三步,这套顺序帮我省下的时间足够再跑一轮完整的消融实验。希望你也能少踩这些坑。
本文还有配套的精品资源,点击获取