基于BiLSTM-CRF和知识图谱的实体识别系统实战：从模型构建到生产部署

1次阅读

共计 2587 个字符，预计需要花费 7 分钟才能阅读完成。

在医疗、金融等专业领域的实体识别任务中，纯序列模型面临两大核心挑战：

专业术语 OOV 问题 ：领域专有名词在通用语料中出现频率低，导致词向量表征不足。例如在医疗文本中，” 二甲双胍 ” 可能被切分为 subword 单元，丢失语义完整性
长距离依赖捕捉困难 ：传统 BiLSTM 在超过 20 个 token 的跨距实体识别时，F1 值平均下降 17.6%。例如金融合同中的 ” 甲方应在签署后 30 个工作日内支付违约金 ”，关键实体间存在复杂语法结构

模型类型	CONLL2003 F1	医疗数据集 Recall	显存占用 (MB)
BiLSTM-CRF	91.2	68.5	1024
BERT-CRF	92.7	75.8	4096
本方案	93.1	82.3	1536

class GatedBiLSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        # 输入门控参数
        self.input_gate = nn.Linear(input_dim + hidden_dim, hidden_dim)
        # 标准 LSTM 参数
        self.lstm = nn.LSTM(input_dim, hidden_dim // 2, bidirectional=True)

    def forward(self, x):
        h_0 = torch.zeros(2, x.size(0), self.hidden_dim // 2).to(x.device)
        c_0 = torch.zeros_like(h_0)

        # 计算门控权重
        gate_input = torch.cat([x, h_0.repeat(1, 1, 2)], dim=-1)
        gate = torch.sigmoid(self.input_gate(gate_input))

        # 应用门控
        x = x * gate
        outputs, _ = self.lstm(x, (h_0, c_0))
        return outputs

使用 PyKEEN 加载 TransE 嵌入：

from pykeen.models import TransE
trans_e = TransE(triples_factory=kg_triples)
entity_embeddings = trans_e.entity_embeddings.weight

字符级 CNN 特征提取器：

class CharCNN(nn.Module):
    def __init__(self, char_vocab_size, embed_dim=50):
        super().__init__()
        self.embed = nn.Embedding(char_vocab_size, embed_dim)
        self.conv = nn.Conv1d(embed_dim, 100, kernel_size=3)

    def forward(self, chars):
        # chars: (batch_size, seq_len, word_len)
        batch_size = chars.size(0)
        x = self.embed(chars)  # (B,S,W,E)
        x = x.view(-1, x.size(2), x.size(3))  # (B*S,W,E)
        x = x.permute(0, 2, 1)  # (B*S,E,W)
        x = F.relu(self.conv(x))  # (B*S,C,W-2)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)  # (B*S,C)
        return x.view(batch_size, -1, 100)  # (B,S,C)

特征拼接层：

class FeatureFusion(nn.Module):
    def __init__(self, lstm_dim, kg_dim, char_dim):
        super().__init__()
        self.proj = nn.Linear(lstm_dim + kg_dim + char_dim, lstm_dim)

    def forward(self, lstm_out, kg_emb, char_feat):
        # kg_emb 需与输入序列对齐
        combined = torch.cat([lstm_out, kg_emb, char_feat], dim=-1)
        return F.gelu(self.proj(combined))

梯度检查点 ：

from torch.utils.checkpoint import checkpoint

def forward_with_checkpoint(x):
    def create_custom_forward(module):
        def custom_forward(*inputs):
            return module(inputs[0])
        return custom_forward

    # 对 BiLSTM 层启用检查点
    x = checkpoint(create_custom_forward(self.bilstm), x)
    return x

混合精度训练 ：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

维度比例建议 ：
图谱嵌入维度 ≤ LSTM 隐藏层的 1 /3
字符 CNN 输出维度建议 50-100

实体嵌套解决方案 ：

# 使用层次化标签方案
LABEL_SCHEME = {
    "B-DISEASE": 0,
    "I-DISEASE": 1,
    "B-DRUG": 2,
    "I-DRUG": 3,
    "B-DISEASE_DRUG": 4,  # 嵌套实体特殊标签
    "I-DISEASE_DRUG": 5
}