用PyTorch复现CasRel关系抽取模型：从百度数据到实战部署的完整流程-平芜编程栈

PyTorch实战：从零构建CasRel模型实现中文关系抽取全流程

当处理"李柏光毕业于北京大学法律系"这样的句子时，传统的关系抽取模型往往难以处理"李柏光-毕业院校-北京大学"和"北京大学-所在地-北京"这类重叠关系。本文将带您完整实现CasRel这一创新性级联二元标注框架，从数据准备到生产部署，解决中文关系抽取中的重叠三元组难题。

1. 环境准备与数据工程

1.1 百度数据集的深度处理

百度开源的关系抽取数据集采用JSONL格式，每条记录包含原始文本和SPO三元组列表。我们需要特别注意中文文本的特殊处理：

{ "text": "《骑士之爱与游吟诗人》是上海社会科学院出版社2012年出版的图书", "spo_list": [ { "predicate": "出版社", "object": "上海社会科学院出版社", "subject": "骑士之爱与游吟诗人" }, { "predicate": "出版时间", "object": "2012年", "subject": "骑士之爱与游吟诗人" } ] }

关键处理步骤：

使用BERT tokenizer时设置add_special_tokens=False保留原始字符映射
对长文本采用动态padding策略，每个batch单独计算max_length
构建subject到(relation, object)的映射字典，处理一对多关系

1.2 高效Batch生成器设计

传统方法逐个处理样本效率低下，我们实现批量化矩阵运算：

class RelationBatchGenerator: def __init__(self, tokenizer, max_len=256): self.tokenizer = tokenizer self.max_len = max_len def __call__(self, batch): texts = [item[0] for item in batch] triples = [item[1] for item in batch] # 动态padding encodings = self.tokenizer( texts, padding=True, truncation=True, max_length=self.max_len, return_tensors="pt" ) # 构建三维标签矩阵 batch_size = len(texts) seq_len = encodings['input_ids'].shape[1] num_rels = len(RELATION_TYPES) sub_head_labels = torch.zeros(batch_size, seq_len) sub_tail_labels = torch.zeros_like(sub_head_labels) obj_head_labels = torch.zeros(batch_size, seq_len, num_rels) obj_tail_labels = torch.zeros_like(obj_head_labels) # 填充标签矩阵 for batch_idx, spo_list in enumerate(triples): # 处理subject标注 for sub_start, sub_end in self._get_entity_spans(spo_list, 'subject'): sub_head_labels[batch_idx, sub_start] = 1 sub_tail_labels[batch_idx, sub_end] = 1 # 处理object标注 for rel, obj_start, obj_end in self._get_relation_spans(spo_list): obj_head_labels[batch_idx, obj_start, rel] = 1 obj_tail_labels[batch_idx, obj_end, rel] = 1 return { 'input_ids': encodings['input_ids'], 'attention_mask': encodings['attention_mask'], 'sub_head_labels': sub_head_labels, 'sub_tail_labels': sub_tail_labels, 'obj_head_labels': obj_head_labels, 'obj_tail_labels': obj_tail_labels }

注意：中文BERT的tokenizer可能导致字符偏移问题，建议保存原始字符位置映射表

2. 模型架构深度解析

2.1 级联标注框架实现

CasRel的核心创新在于将关系抽取分解为两个子任务：

Subject识别模块：采用标准的序列标注
Relation-specific对象识别：针对每个候选subject预测可能的关系和对象

class CasRelModel(nn.Module): def __init__(self, pretrained_path, num_relations): super().__init__() self.bert = BertModel.from_pretrained(pretrained_path) hidden_size = self.bert.config.hidden_size # Subject识别头 self.sub_head_linear = nn.Linear(hidden_size, 1) self.sub_tail_linear = nn.Linear(hidden_size, 1) # Relation-specific对象识别头 self.obj_head_linear = nn.Linear(hidden_size, num_relations) self.obj_tail_linear = nn.Linear(hidden_size, num_relations) def forward(self, input_ids, attention_mask): outputs = self.bert(input_ids, attention_mask=attention_mask) sequence_output = outputs.last_hidden_state # Subject预测 sub_head_logits = torch.sigmoid(self.sub_head_linear(sequence_output)) sub_tail_logits = torch.sigmoid(self.sub_tail_linear(sequence_output)) # 对每个候选subject预测对象 batch_size, seq_len, _ = sequence_output.shape sub_head_mask = (sub_head_logits > 0.5).float() sub_rep = torch.bmm(sub_head_mask.transpose(1,2), sequence_output) # 加入subject信息增强表示 enhanced_output = sequence_output + sub_rep # 对象预测 obj_head_logits = torch.sigmoid(self.obj_head_linear(enhanced_output)) obj_tail_logits = torch.sigmoid(self.obj_tail_linear(enhanced_output)) return { 'sub_head': sub_head_logits, 'sub_tail': sub_tail_logits, 'obj_head': obj_head_logits, 'obj_tail': obj_tail_logits }

2.2 焦点损失函数优化

针对正负样本不平衡问题，我们改进原始二元交叉熵损失：

class FocalLoss(nn.Module): def __init__(self, alpha=0.25, gamma=2): super().__init__() self.alpha = alpha self.gamma = gamma def forward(self, preds, targets, mask): # 过滤padding部分 preds = preds[mask == 1] targets = targets[mask == 1] bce_loss = F.binary_cross_entropy(preds, targets, reduction='none') p_t = torch.where(targets == 1, preds, 1-preds) alpha_t = torch.where(targets == 1, self.alpha, 1-self.alpha) loss = alpha_t * (1 - p_t)**self.gamma * bce_loss return loss.mean() # 组合损失 def compute_loss(model_output, labels, mask): sub_head_loss = focal_loss(model_output['sub_head'], labels['sub_head'], mask) sub_tail_loss = focal_loss(model_output['sub_tail'], labels['sub_tail'], mask) # 对象损失需要考虑所有关系类型 obj_mask = mask.unsqueeze(-1).expand(-1, -1, labels['obj_head'].shape[-1]) obj_head_loss = focal_loss(model_output['obj_head'], labels['obj_head'], obj_mask) obj_tail_loss = focal_loss(model_output['obj_tail'], labels['obj_tail'], obj_mask) return sub_head_loss + sub_tail_loss + obj_head_loss + obj_tail_loss

3. 训练策略与技巧

3.1 渐进式学习率调度

采用线性warmup和余弦退火组合策略：

def get_optimizer(model, lr, warmup_steps, total_steps): no_decay = ['bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [ { 'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01 }, { 'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 } ] optimizer = AdamW(optimizer_grouped_parameters, lr=lr, eps=1e-8) scheduler = get_cosine_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps ) return optimizer, scheduler

3.2 对抗训练增强

通过FGM对抗训练提升模型鲁棒性：

class FGM: def __init__(self, model): self.model = model self.backup = {} def attack(self, epsilon=0.5): for name, param in self.model.named_parameters(): if param.requires_grad and param.grad is not None: self.backup[name] = param.data.clone() norm = torch.norm(param.grad) if norm != 0: r_at = epsilon * param.grad / norm param.data.add_(r_at) def restore(self): for name, param in self.model.named_parameters(): if name in self.backup: param.data = self.backup[name] self.backup = {} # 训练循环中使用 fgm = FGM(model) loss.backward() fgm.attack() # 在梯度上施加扰动 loss_adv = compute_loss(model(input_ids, attention_mask), labels) loss_adv.backward() fgm.restore() optimizer.step()

4. 部署优化与性能调优

4.1 ONNX运行时加速

将PyTorch模型导出为ONNX格式：

def export_onnx(model, output_path): dummy_input = { 'input_ids': torch.randint(0, 100, (1, 128)), 'attention_mask': torch.ones((1, 128)) } torch.onnx.export( model, (dummy_input['input_ids'], dummy_input['attention_mask']), output_path, input_names=['input_ids', 'attention_mask'], output_names=['sub_head', 'sub_tail', 'obj_head', 'obj_tail'], dynamic_axes={ 'input_ids': {0: 'batch', 1: 'sequence'}, 'attention_mask': {0: 'batch', 1: 'sequence'}, 'sub_head': {0: 'batch', 1: 'sequence'}, 'sub_tail': {0: 'batch', 1: 'sequence'}, 'obj_head': {0: 'batch', 1: 'sequence', 2: 'relations'}, 'obj_tail': {0: 'batch', 1: 'sequence', 2: 'relations'} }, opset_version=12 )

4.2 Triton推理服务器部署

配置Triton的config.pbtxt：

name: "casrel_relation_extraction" platform: "onnxruntime_onnx" max_batch_size: 32 input [ { name: "input_ids" data_type: TYPE_INT64 dims: [ -1, -1 ] }, { name: "attention_mask" data_type: TYPE_INT64 dims: [ -1, -1 ] } ] output [ { name: "sub_head" data_type: TYPE_FP32 dims: [ -1, -1 ] }, { name: "sub_tail" data_type: TYPE_FP32 dims: [ -1, -1 ] }, { name: "obj_head" data_type: TYPE_FP32 dims: [ -1, -1, -1 ] }, { name: "obj_tail" data_type: TYPE_FP32 dims: [ -1, -1, -1 ] } ]