从翻译到旋转用Python代码复现TransE到RotatE的演进之路知识表示学习是知识图谱领域的核心技术之一它将实体和关系映射到连续向量空间使得机器能够更好地理解和推理知识。本文将带您从零开始用Python代码实现从TransE到RotatE的完整演进过程通过实践理解这些经典模型的设计哲学与技术细节。1. 知识表示学习基础环境搭建在开始实现之前我们需要配置开发环境。推荐使用Python 3.8和PyTorch 1.10这些版本在性能和兼容性方面都有良好表现。# 创建虚拟环境 python -m venv kg_env source kg_env/bin/activate # Linux/Mac kg_env\Scripts\activate # Windows # 安装核心依赖 pip install torch1.13.1 numpy1.23.5 matplotlib3.6.2提示如果使用GPU加速训练请安装对应CUDA版本的PyTorch知识图谱数据通常以三元组形式存储。我们定义一个简单的数据加载器class KGDataset: def __init__(self, triples): self.entities set() self.relations set() for h, r, t in triples: self.entities.update([h, t]) self.relations.add(r) self.entity2id {e: i for i, e in enumerate(self.entities)} self.relation2id {r: i for i, r in enumerate(self.relations)} self.triples [(self.entity2id[h], self.relation2id[r], self.entity2id[t]) for h, r, t in triples]2. TransE模型实现与优化TransE作为知识表示学习的开山之作其核心思想是将关系视为头实体到尾实体的平移向量。我们先实现基础版本import torch import torch.nn as nn import torch.nn.functional as F class TransE(nn.Module): def __init__(self, num_entities, num_relations, embedding_dim50): super(TransE, self).__init__() self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim) # 初始化参数 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.xavier_uniform_(self.relation_emb.weight) def forward(self, h, r, t): h_emb self.entity_emb(h) r_emb self.relation_emb(r) t_emb self.entity_emb(t) # 计算得分函数 score torch.norm(h_emb r_emb - t_emb, p2, dim1) return score训练TransE需要特别设计负采样策略和损失函数def train_transE(model, dataset, epochs100, lr0.01, margin1.0): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0 for h, r, t in dataset.triples: # 正样本 pos_score model(h, r, t) # 负采样随机替换头或尾实体 if torch.rand(1) 0.5: neg_h torch.randint(0, len(dataset.entities), (1,)) neg_score model(neg_h, r, t) else: neg_t torch.randint(0, len(dataset.entities), (1,)) neg_score model(h, r, neg_t) # 最大化正负样本得分差距 loss F.relu(margin pos_score - neg_score) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataset.triples):.4f})TransE虽然简单高效但在处理复杂关系时存在明显局限。下表展示了不同关系类型下的表现差异关系类型FB15k-237准确率WN18RR准确率1-to-10.780.851-to-N0.620.58N-to-10.590.53N-to-N0.510.473. TransH模型超平面投影的改进TransH通过将实体投影到关系特定的超平面来解决复杂关系问题。以下是关键实现class TransH(nn.Module): def __init__(self, num_entities, num_relations, embedding_dim50): super(TransH, self).__init__() self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim) self.norm_vectors nn.Embedding(num_relations, embedding_dim) # 初始化参数 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.xavier_uniform_(self.relation_emb.weight) nn.init.xavier_uniform_(self.norm_vectors.weight) def project(self, e, norm): norm F.normalize(norm, p2, dim-1) return e - torch.sum(e * norm, dim-1, keepdimTrue) * norm def forward(self, h, r, t): h_emb self.entity_emb(h) r_emb self.relation_emb(r) t_emb self.entity_emb(t) norm self.norm_vectors(r) # 投影到超平面 h_proj self.project(h_emb, norm) t_proj self.project(t_emb, norm) score torch.norm(h_proj r_emb - t_proj, p2, dim1) return scoreTransH的训练过程与TransE类似但需要特别注意超平面法向量的归一化def train_transH(model, dataset, epochs100, lr0.01, margin1.0): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0 for h, r, t in dataset.triples: # 正样本得分 pos_score model(h, r, t) # 负采样 if torch.rand(1) 0.5: neg_h torch.randint(0, len(dataset.entities), (1,)) neg_score model(neg_h, r, t) else: neg_t torch.randint(0, len(dataset.entities), (1,)) neg_score model(h, r, neg_t) loss F.relu(margin pos_score - neg_score) optimizer.zero_grad() loss.backward() # 手动归一化法向量 model.norm_vectors.weight.data F.normalize( model.norm_vectors.weight.data, p2, dim-1) optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataset.triples):.4f})4. RotatE模型复数空间中的旋转RotatE将关系建模为复数空间中的旋转操作这是知识表示学习的重要突破。我们首先需要实现复数运算def complex_mul(h, r): 复数乘法实现旋转操作 h_re, h_im torch.chunk(h, 2, dim-1) r_re, r_im torch.chunk(r, 2, dim-1) return torch.cat([ h_re * r_re - h_im * r_im, h_re * r_im h_im * r_re ], dim-1) class RotatE(nn.Module): def __init__(self, num_entities, num_relations, embedding_dim50): super(RotatE, self).__init__() # 复数维度是实际维度的一半 self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim//2) # 初始化参数 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.uniform_(self.relation_emb.weight, a0, b2*torch.pi) def forward(self, h, r, t): h_emb self.entity_emb(h) r_phase self.relation_emb(r) t_emb self.entity_emb(t) # 将关系转换为复数形式 r_emb torch.cat([torch.cos(r_phase), torch.sin(r_phase)], dim-1) # 旋转操作 h_rot complex_mul(h_emb, r_emb) # 计算距离 score torch.norm(h_rot - t_emb, p2, dim1) return scoreRotatE采用了不同的损失函数设计借鉴了负采样和自对抗训练的思想def train_rotate(model, dataset, epochs100, lr0.01, margin1.0, neg_sample_size10): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0 for h, r, t in dataset.triples: # 正样本得分 pos_score model(h, r, t) # 负采样 neg_scores [] for _ in range(neg_sample_size): if torch.rand(1) 0.5: neg_h torch.randint(0, len(dataset.entities), (1,)) neg_score model(neg_h, r, t) else: neg_t torch.randint(0, len(dataset.entities), (1,)) neg_score model(h, r, neg_t) neg_scores.append(neg_score) neg_scores torch.stack(neg_scores) # 自对抗负采样权重 weights F.softmax(neg_scores.detach(), dim0) # 加权负采样损失 loss -F.logsigmoid(margin - pos_score) - \ torch.sum(weights * F.logsigmoid(neg_scores - margin)) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataset.triples):.4f})5. 模型评估与可视化分析实现模型后我们需要评估其性能并可视化嵌入结果。常用的评估指标包括Mean Rank (MR): 正确三元组排名的平均值Mean Reciprocal Rank (MRR): 排名倒数的平均值HitsK: 正确三元组排名在前K的比例def evaluate(model, dataset, test_triples): ranks [] for h, r, t in test_triples: # 计算所有尾实体的得分 all_t torch.arange(len(dataset.entities)) h_batch torch.full_like(all_t, h) r_batch torch.full_like(all_t, r) scores model(h_batch, r_batch, all_t) # 获取正确尾实体的排名 sorted_indices torch.argsort(scores) rank (sorted_indices t).nonzero().item() 1 ranks.append(rank) mr torch.tensor(ranks).float().mean().item() mrr (1. / torch.tensor(ranks).float()).mean().item() hits10 (torch.tensor(ranks) 10).float().mean().item() return {MR: mr, MRR: mrr, Hits10: hits10}可视化嵌入结果可以帮助我们直观理解模型学习到的模式import matplotlib.pyplot as plt from sklearn.decomposition import PCA def visualize_embeddings(model, dataset, relation_id): # 获取实体和关系嵌入 entity_emb model.entity_emb.weight.detach().numpy() relation_emb model.relation_emb.weight.detach().numpy() # 使用PCA降维 pca PCA(n_components2) entity_2d pca.fit_transform(entity_emb) # 可视化 plt.figure(figsize(10, 8)) plt.scatter(entity_2d[:, 0], entity_2d[:, 1], alpha0.5) # 标记特定关系 related_triples [t for t in dataset.triples if t[1] relation_id] for h, r, t in related_triples[:20]: # 只显示部分 plt.plot([entity_2d[h, 0], entity_2d[t, 0]], [entity_2d[h, 1], entity_2d[t, 1]], r-, alpha0.3) plt.text(entity_2d[h, 0], entity_2d[h, 1], str(h), fontsize8) plt.text(entity_2d[t, 0], entity_2d[t, 1], str(t), fontsize8) plt.title(fEmbedding Visualization for Relation {relation_id}) plt.show()在实际项目中我发现RotatE的复数旋转操作虽然数学上优雅但在实现时需要特别注意数值稳定性。特别是在计算相位角时加入小的epsilon值可以避免NaN问题# 改进的RotatE关系嵌入初始化 self.relation_emb nn.Embedding(num_relations, embedding_dim//2) nn.init.uniform_(self.relation_emb.weight, a-0.01, b0.01) # 小范围初始化 # 前向传播中加入稳定性处理 r_phase torch.clamp(self.relation_emb(r), -math.pi, math.pi) r_emb torch.cat([ torch.cos(r_phase) 1e-10, torch.sin(r_phase) 1e-10 ], dim-1)
从翻译到旋转：用Python代码复现TransE到RotatE的演进之路（附GitHub源码）
发布时间：2026/5/20 17:46:45
从翻译到旋转用Python代码复现TransE到RotatE的演进之路知识表示学习是知识图谱领域的核心技术之一它将实体和关系映射到连续向量空间使得机器能够更好地理解和推理知识。本文将带您从零开始用Python代码实现从TransE到RotatE的完整演进过程通过实践理解这些经典模型的设计哲学与技术细节。1. 知识表示学习基础环境搭建在开始实现之前我们需要配置开发环境。推荐使用Python 3.8和PyTorch 1.10这些版本在性能和兼容性方面都有良好表现。# 创建虚拟环境 python -m venv kg_env source kg_env/bin/activate # Linux/Mac kg_env\Scripts\activate # Windows # 安装核心依赖 pip install torch1.13.1 numpy1.23.5 matplotlib3.6.2提示如果使用GPU加速训练请安装对应CUDA版本的PyTorch知识图谱数据通常以三元组形式存储。我们定义一个简单的数据加载器class KGDataset: def __init__(self, triples): self.entities set() self.relations set() for h, r, t in triples: self.entities.update([h, t]) self.relations.add(r) self.entity2id {e: i for i, e in enumerate(self.entities)} self.relation2id {r: i for i, r in enumerate(self.relations)} self.triples [(self.entity2id[h], self.relation2id[r], self.entity2id[t]) for h, r, t in triples]2. TransE模型实现与优化TransE作为知识表示学习的开山之作其核心思想是将关系视为头实体到尾实体的平移向量。我们先实现基础版本import torch import torch.nn as nn import torch.nn.functional as F class TransE(nn.Module): def __init__(self, num_entities, num_relations, embedding_dim50): super(TransE, self).__init__() self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim) # 初始化参数 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.xavier_uniform_(self.relation_emb.weight) def forward(self, h, r, t): h_emb self.entity_emb(h) r_emb self.relation_emb(r) t_emb self.entity_emb(t) # 计算得分函数 score torch.norm(h_emb r_emb - t_emb, p2, dim1) return score训练TransE需要特别设计负采样策略和损失函数def train_transE(model, dataset, epochs100, lr0.01, margin1.0): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0 for h, r, t in dataset.triples: # 正样本 pos_score model(h, r, t) # 负采样随机替换头或尾实体 if torch.rand(1) 0.5: neg_h torch.randint(0, len(dataset.entities), (1,)) neg_score model(neg_h, r, t) else: neg_t torch.randint(0, len(dataset.entities), (1,)) neg_score model(h, r, neg_t) # 最大化正负样本得分差距 loss F.relu(margin pos_score - neg_score) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataset.triples):.4f})TransE虽然简单高效但在处理复杂关系时存在明显局限。下表展示了不同关系类型下的表现差异关系类型FB15k-237准确率WN18RR准确率1-to-10.780.851-to-N0.620.58N-to-10.590.53N-to-N0.510.473. TransH模型超平面投影的改进TransH通过将实体投影到关系特定的超平面来解决复杂关系问题。以下是关键实现class TransH(nn.Module): def __init__(self, num_entities, num_relations, embedding_dim50): super(TransH, self).__init__() self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim) self.norm_vectors nn.Embedding(num_relations, embedding_dim) # 初始化参数 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.xavier_uniform_(self.relation_emb.weight) nn.init.xavier_uniform_(self.norm_vectors.weight) def project(self, e, norm): norm F.normalize(norm, p2, dim-1) return e - torch.sum(e * norm, dim-1, keepdimTrue) * norm def forward(self, h, r, t): h_emb self.entity_emb(h) r_emb self.relation_emb(r) t_emb self.entity_emb(t) norm self.norm_vectors(r) # 投影到超平面 h_proj self.project(h_emb, norm) t_proj self.project(t_emb, norm) score torch.norm(h_proj r_emb - t_proj, p2, dim1) return scoreTransH的训练过程与TransE类似但需要特别注意超平面法向量的归一化def train_transH(model, dataset, epochs100, lr0.01, margin1.0): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0 for h, r, t in dataset.triples: # 正样本得分 pos_score model(h, r, t) # 负采样 if torch.rand(1) 0.5: neg_h torch.randint(0, len(dataset.entities), (1,)) neg_score model(neg_h, r, t) else: neg_t torch.randint(0, len(dataset.entities), (1,)) neg_score model(h, r, neg_t) loss F.relu(margin pos_score - neg_score) optimizer.zero_grad() loss.backward() # 手动归一化法向量 model.norm_vectors.weight.data F.normalize( model.norm_vectors.weight.data, p2, dim-1) optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataset.triples):.4f})4. RotatE模型复数空间中的旋转RotatE将关系建模为复数空间中的旋转操作这是知识表示学习的重要突破。我们首先需要实现复数运算def complex_mul(h, r): 复数乘法实现旋转操作 h_re, h_im torch.chunk(h, 2, dim-1) r_re, r_im torch.chunk(r, 2, dim-1) return torch.cat([ h_re * r_re - h_im * r_im, h_re * r_im h_im * r_re ], dim-1) class RotatE(nn.Module): def __init__(self, num_entities, num_relations, embedding_dim50): super(RotatE, self).__init__() # 复数维度是实际维度的一半 self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim//2) # 初始化参数 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.uniform_(self.relation_emb.weight, a0, b2*torch.pi) def forward(self, h, r, t): h_emb self.entity_emb(h) r_phase self.relation_emb(r) t_emb self.entity_emb(t) # 将关系转换为复数形式 r_emb torch.cat([torch.cos(r_phase), torch.sin(r_phase)], dim-1) # 旋转操作 h_rot complex_mul(h_emb, r_emb) # 计算距离 score torch.norm(h_rot - t_emb, p2, dim1) return scoreRotatE采用了不同的损失函数设计借鉴了负采样和自对抗训练的思想def train_rotate(model, dataset, epochs100, lr0.01, margin1.0, neg_sample_size10): optimizer torch.optim.Adam(model.parameters(), lrlr) for epoch in range(epochs): total_loss 0 for h, r, t in dataset.triples: # 正样本得分 pos_score model(h, r, t) # 负采样 neg_scores [] for _ in range(neg_sample_size): if torch.rand(1) 0.5: neg_h torch.randint(0, len(dataset.entities), (1,)) neg_score model(neg_h, r, t) else: neg_t torch.randint(0, len(dataset.entities), (1,)) neg_score model(h, r, neg_t) neg_scores.append(neg_score) neg_scores torch.stack(neg_scores) # 自对抗负采样权重 weights F.softmax(neg_scores.detach(), dim0) # 加权负采样损失 loss -F.logsigmoid(margin - pos_score) - \ torch.sum(weights * F.logsigmoid(neg_scores - margin)) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataset.triples):.4f})5. 模型评估与可视化分析实现模型后我们需要评估其性能并可视化嵌入结果。常用的评估指标包括Mean Rank (MR): 正确三元组排名的平均值Mean Reciprocal Rank (MRR): 排名倒数的平均值HitsK: 正确三元组排名在前K的比例def evaluate(model, dataset, test_triples): ranks [] for h, r, t in test_triples: # 计算所有尾实体的得分 all_t torch.arange(len(dataset.entities)) h_batch torch.full_like(all_t, h) r_batch torch.full_like(all_t, r) scores model(h_batch, r_batch, all_t) # 获取正确尾实体的排名 sorted_indices torch.argsort(scores) rank (sorted_indices t).nonzero().item() 1 ranks.append(rank) mr torch.tensor(ranks).float().mean().item() mrr (1. / torch.tensor(ranks).float()).mean().item() hits10 (torch.tensor(ranks) 10).float().mean().item() return {MR: mr, MRR: mrr, Hits10: hits10}可视化嵌入结果可以帮助我们直观理解模型学习到的模式import matplotlib.pyplot as plt from sklearn.decomposition import PCA def visualize_embeddings(model, dataset, relation_id): # 获取实体和关系嵌入 entity_emb model.entity_emb.weight.detach().numpy() relation_emb model.relation_emb.weight.detach().numpy() # 使用PCA降维 pca PCA(n_components2) entity_2d pca.fit_transform(entity_emb) # 可视化 plt.figure(figsize(10, 8)) plt.scatter(entity_2d[:, 0], entity_2d[:, 1], alpha0.5) # 标记特定关系 related_triples [t for t in dataset.triples if t[1] relation_id] for h, r, t in related_triples[:20]: # 只显示部分 plt.plot([entity_2d[h, 0], entity_2d[t, 0]], [entity_2d[h, 1], entity_2d[t, 1]], r-, alpha0.3) plt.text(entity_2d[h, 0], entity_2d[h, 1], str(h), fontsize8) plt.text(entity_2d[t, 0], entity_2d[t, 1], str(t), fontsize8) plt.title(fEmbedding Visualization for Relation {relation_id}) plt.show()在实际项目中我发现RotatE的复数旋转操作虽然数学上优雅但在实现时需要特别注意数值稳定性。特别是在计算相位角时加入小的epsilon值可以避免NaN问题# 改进的RotatE关系嵌入初始化 self.relation_emb nn.Embedding(num_relations, embedding_dim//2) nn.init.uniform_(self.relation_emb.weight, a-0.01, b0.01) # 小范围初始化 # 前向传播中加入稳定性处理 r_phase torch.clamp(self.relation_emb(r), -math.pi, math.pi) r_emb torch.cat([ torch.cos(r_phase) 1e-10, torch.sin(r_phase) 1e-10 ], dim-1)
