简介这是一份面向知识图谱、自然语言处理与机器学习学习者的TransE模型Python实现基于FB15k数据集进行训练适合想通过代码理解知识图谱嵌入原理、三元组关系建模hr≈t的读者。压缩包共21个文件大小5.85MB包括4个Python源码、14个txt数据与说明文件、2个md文档和1个png示意图txt数据中包含了train.txt的训练集划分py脚本覆盖模型实现md文件提供讲解说明png展示模型结构整体目录结构清晰便于按需查阅。已有403人学习下载。借助这套代码可以学习实体与关系向量的初始化、损失函数设计、负采样策略等关键步骤并能在FB15k数据上完成链接预测、知识图谱补全等任务md文档与数据文件的搭配也有助于从数据处理到模型评估的完整复盘适合初学者进阶或作为课程实验参考。1. 拿到“transE 模型 Python 版实现 FB15k 数据集”这个包先别急着训练网上以 transE 模型 python 版实现、FB15k 数据集、train.txt 为核心词的压缩包很多下载之后最常出现的尴尬是跑通了 loss 也在降最后 mean rank 却高得离谱连自己都不知道模型到底学没学到东西。问题一般不在模型本身而在数据拆分这一环。TransE 要处理的 FB15k 是知识图谱三元组数据的标准子集每一行只有头实体、关系、尾实体三个词不对它做负采样就直接喂网络模型会退化到把所有向量挤压到一起。这篇文章把一条可复现的路径拆开讲train.txt 怎么读、负样本怎么拆、参数怎么调、评测哪里最容易翻车。适合想跑通知识图谱嵌入做链接预测、或准备复现论文指标的工程师和学生。2. FB15k 数据集与 train.txt 解析三列文本、实体映射和文件划分2.1 先认清 FB15k 的三件套train.txt、valid.txt、test.txt 各自做什么FB15k 是 Freebase 上截出来的知识图谱子集规模不大不小实体一万多、关系一千多正好适合在单机上做 TransE 实验。一个标准的资源包解压后你通常会看到 train.txt、valid.txt、test.txt 三个文本文件配套的还有 entity2id.txt 和 relation2id.txt。前三个文件存的是三元组后两个文件存的是名字到数字 id 的映射。每个 txt 里一行一个三元组最常见分隔符是 tab偶尔有空格分隔的版本。train.txt 的典型格式是“实体名\t关系名\t实体名”比如某一行可能长成/m/02mjmr /film/film/featured_film /m/02mjmr实际上文件里就是三列字符串。新手最容易犯的第一个错是把这三个文件理解成“训练集、验证集、测试集”之后自己重新划分数据。其实 TransE 的训练评估早就有固定分工train.txt 负责提供正样本valid.txt 负责做超参挑选和早期停止test.txt 只在完全训练完之后碰一次。我见过不少人在拿到压缩包之后先把 train.txt 里 90% 的行当训练集、10% 当验证集然后把 valid.txt 和 test.txt 丢在一边。这个动作不是不行但会破坏和论文指标的对比基础。FB15k 的划分本身是固定的train/valid/test 三个集合互不重叠你直接按文件分工使用后面复现别人结果时才能对上号。数据体检这一步很重要先搞清楚手里这份 zip 里的 train.txt 究竟长什么样再决定要不要动它。2.2 用 Python 把 train.txt 读成三元组并顺手建立 entity2id 映射解压之后不要先急着写模型先把数据读进来用 Python 做一次数据体检。读样本的脚本看起来简单但分隔符和脏行往往在第一个晚上就把你卡住。我一般用下面这段兼容 tab 和空格的处理def load_triples(path): triples [] with open(path, r, encodingutf-8) as f: for line in f: line line.strip() if not line: continue parts line.split(\t) if len(parts) ! 3: parts line.split() if len(parts) ! 3: continue # 跳过空行、注释行或明显损坏的行 h, r, t parts triples.append((h, r, t)) return triples train_triples load_triples(train.txt) valid_triples load_triples(valid.txt) test_triples load_triples(test.txt) print(train:, len(train_triples), valid:, len(valid_triples), test:, len(test_triples)) print(样例三元组:, train_triples[0])这段代码的逻辑很直白先去掉每行首尾空白空行直接跳过优先按 tab 切分如果切不出三列再按空格切一次还不行就丢弃。这样即使下载到的包是空格版本也不会读挂。接下来要做的事是给实体和关系分别建 id。注意这里的关键点映射一定要用 train、valid、test 三个文件的全部数据一起建不能用完 train.txt 再单独建一次。否则测试集里出现训练集里没有的实体时你后面转 id 会直接 KeyError或者在 debug 时留下一堆 never seen 实体。def build_mappings(all_triples): entity2id {} relation2id {} for h, r, t in all_triples: if h not in entity2id: entity2id[h] len(entity2id) if t not in entity2id: entity2id[t] len(entity2id) if r not in relation2id: relation2id[r] len(relation2id) return entity2id, relation2id all_text_triples train_triples valid_triples test_triples entity2id, relation2id build_mappings(all_text_triples) print(实体总数:, len(entity2id), 关系总数:, len(relation2id))这里我用 dict 的 len 作为自增 id而不是猜一个数字范围原因是 FB15k 的原始实体名是一长串 URI 风格字符串用 len 当序号可以保证 id 连续且无冲突。id 从 0 开始还是从 1 开始无所谓但后面用 PyTorch 的 Embedding 层时id 最大绝对不能超过 num_embeddings 减 1否则 IndexError。我见过有人从 1 开始编号结果 embedding 层多申请了一个向量训练时还把 0 号向量空置这种细节不影响结果但会在保存映射表时把自己绕晕。文本三元组读进来之后建议把它们转成纯数字的三元组数组训练时只吃数字数组。这样后面负采样和 token 化都省事train_ids [] for h, r, t in train_triples: train_ids.append((entity2id[h], relation2id[r], entity2id[t]))加上这个转换后训练过程就是操作 int不用反复查 dict内存占用也更小。与符号名的耦合点收敛到 entity2id/relation2id 这两个映射上后面保存模型时把这两个映射一起存好就是给自己留了后悔药。3. 数据拆分的本质是负采样为什么正样本三元组不能直接喂给 TransE3.1 TransE 的翻译假设与损失函数距离公式和 margin ranking loss 怎么对应TransE 的基本假设是“翻译”对于一个正确的三元组 (h, r, t)在向量空间里应该满足 h r ≈ t也就是实体 h 沿关系 r 的方向平移一段距离就能到达实体 t。用距离来度量是score ||h r - t||距离越小代表三元组成立的可能性越大。如果你只把 train.txt 里的正样本喂进去模型会找到一个非常取巧的解把所有向量初始化成同一个点再把关系向量也设成零向量这样任意三元组的距离都是 0loss 直接归零但模型什么也没学到。为了逼模型把真正的结构信息编码进向量必须给每个正样本配一个负样本让模型看到“把 h 或 t 换掉之后距离应该变大”。对应的损失函数是 margin ranking lossloss max(0, margin d_pos - d_neg)。意思是正样本距离 d_pos 至少要比负样本距离 d_neg 小 margin 那么多如果达不到这个间隔就要产生 loss 并更新参数。margin 越大模型被要求拉开的间隔越宽但训练也越难收敛margin 设成 0.5 到 2.0 之间是比较常见的范围。这个阶段最需要注意的是“数据拆分”——不把 train.txt 拆成“模型能看到的正样本”和“用来逼模型区分的负样本”训练就不可能正常。负样本的数量和正样本一一对应常见做法是每个正三元组生成一个 corrupt 三元组。后面如果要做更严格的实验可以再调整负采样比例但第一次跑通用 1:1 就够。3.2 负采样实现随机替换头实体或尾实体以及伪负样本过滤负样本的标准生成方式是随机替换头实体或尾实体。替换哪个没有绝对标准很多实现用 0.5 的概率随机决定有的实现会统计每个关系的平均头/尾实体数量然后用伯努利分布决定替换概率这种方案叫伯努利采样在小数据上提升有限入门阶段直接用 0.5 没问题。下面这段负采样代码接收一个正样本 batch对每条正样本生成一条负样本并且会过滤掉那些已经在 train.txt 里出现过的“伪负样本”import random def generate_negative_batch(batch_triples, entity_ids, train_set, corrupt_prob0.5): neg_batch [] for h, r, t in batch_triples: for _ in range(20): if random.random() corrupt_prob: neg_h random.choice(entity_ids) candidate (neg_h, r, t) else: neg_t random.choice(entity_ids) candidate (h, r, neg_t) if candidate not in train_set: neg_batch.append(candidate) break else: # 20 次都没抽到合法负样本说明这个三元组周围太密 # 退回正样本本身靠 margin 机制把它压下去 neg_batch.append((h, r, t)) return neg_batch # 用 id 三元组构建 set做 O(1) 成员判断 train_set set(train_ids) entity_ids list(range(len(entity2id))) batch train_ids[:128] neg_batch generate_negative_batch(batch, entity_ids, train_set) print(正样本量:, len(batch), 负样本量:, len(neg_batch))几个细节值得说清楚。train_set 必须用 id 三元组而不是文本三元组否则每次 candidate not in train_set 都要比较三个字符串速度会慢十几倍。一个三元组附近密度特别高时比如头实体是某个超大连接节点随机抽 20 次都可能落在真实三元组上代码会退回正样本本身。这种时候不要死循环把它当成负样本用损失函数会让它作为负样本“示弱”只是梯度里有少量噪声不会造成灾难性后果。负采样完成后数据维度是正样本和负样本一一对应训练时两个 batch 一起喂进模型。这里有一个很隐蔽的坑过滤负样本不能只在数据预处理阶段做一次因为训练过程中模型参数还没定型判断“是否在 train_set 中出现”只需要判断三元组是否存在与参数无关所以这个 set 可以持久化复用不需要每次都重新生成。你在重复实验时可以直接 set 固化省去每次采样前的重建时间。4. 训练循环与评估从 batch 到 valid 集上的 mean rank 和 hit104.1 模型实现与 batch 训练embedding 初始化、L2 归一化和可复现脚本TransE 的 PyTorch 实现骨架不复杂两个 embedding 层分别放实体和关系前向计算正负样本的得分差。但有几个细节会影响收敛质量最容易踩的是实体向量和关系向量要不要归一化。原论文里每轮更新后会把实体向量做 L2 归一化这样是为了防止向量模长无限增长实际实现里我在前向时对 embedding 向量先 normalize效果等同且实现更干净。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_dim100, margin1.0): super().__init__() self.entity_emb nn.Embedding(num_entities, embedding_dim) self.relation_emb nn.Embedding(num_relations, embedding_dim) self.margin margin # 常见初始化方式Xavier 均匀分布 nn.init.xavier_uniform_(self.entity_emb.weight) nn.init.xavier_uniform_(self.relation_emb.weight) def forward(self, pos_triples, neg_triples): # pos_triples / neg_triples: [batch, 3] 的 int64 张量 pos_h F.normalize(self.entity_emb(pos_triples[:, 0]), p2, dim1) pos_r F.normalize(self.relation_emb(pos_triples[:, 1]), p2, dim1) pos_t F.normalize(self.entity_emb(pos_triples[:, 2]), p2, dim1) neg_h F.normalize(self.entity_emb(neg_triples[:, 0]), p2, dim1) neg_r F.normalize(self.relation_emb(neg_triples[:, 1]), p2, dim1) neg_t F.normalize(self.entity_emb(neg_triples[:, 2]), p2, dim1) pos_score torch.norm(pos_h pos_r - pos_t, p2, dim1) neg_score torch.norm(neg_h neg_r - neg_t, p2, dim1) loss torch.relu(self.margin pos_score - neg_score).mean() return loss前向函数里把 6 个 embedding 向量全部做 L2 归一化这里有个参数选择关系向量要不要归一化其实可以实验。原论文里对实体向量归一化更关键关系向量归一化后会让关系向量的模长固定好处是训练稳定坏处是损失了一部分表达力。我一般两者都归一化先用默认设置跑到指标可以复现再慢慢放开。训练循环采用 batch 负采样每个 batch 内正负样本数量相等用 Adam 优化器更新参数device torch.device(cuda if torch.cuda.is_available() else cpu) model TransE(len(entity2id), len(relation2id), embedding_dim100, margin1.0).to(device) optimizer torch.optim.Adam(model.parameters(), lr0.001) epochs 500 batch_size 256 for epoch in range(epochs): random.shuffle(train_ids) total_loss 0.0 for i in range(0, len(train_ids), batch_size): batch train_ids[i:i batch_size] neg_batch generate_negative_batch(batch, entity_ids, train_set) pos_tensor torch.tensor(batch, dtypetorch.long, devicedevice) neg_tensor torch.tensor(neg_batch, dtypetorch.long, devicedevice) loss model(pos_tensor, neg_tensor) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() if (epoch 1) % 10 0: print(fepoch {epoch 1}/{epochs}, loss {total_loss / (len(train_ids) / batch_size):.4f})训练循环里值得注意的点是 shuffle 和 batch 划分。FB15k 的 train.txt 原始顺序往往按主题聚在一起如果不 shuffle连续几个 batch 可能只覆盖到少数实体负采样也会随之偏移。另一个点是负采样放在 batch 内部生成这样正样本和负样本之间的对应关系在一次迭代内是稳定的比预先离线生成全部负样本再训练更省内存。4.2 三个必调参数embedding 维度、margin 和学习率附参数对照表第一次跑 TransE不需要把参数调到论文最优先把复现链路跑通然后按下面这张表去感知参数的影响参数常见取值我的一般调法embedding_dim50 / 100 / 200FB15k 默认 100维度增大提升上限但收敛慢、内存涨margin0.5 / 1.0 / 2.0先设 1.0loss 始终不下降再降到 0.5learning_rate0.001 / 0.01Adam 用 0.001SGD 用 0.01 起步batch_size128 / 256 / 512单机 CPU 用 128GPU 用 256 以上epochs300 / 500 / 1000先跑 300 看 valid 指标平稳后再加 200embedding 维度和 FB15k 的关系是实体总数超过一万关系总数约一千维度 100 已经足够表达大部分结构信息维度加到 200 不一定带来 mean rank 的显著下降但训练耗时可能翻倍。margin 是个非常敏感的参数它过大时模型会被要求把正负样本强行拉开训练后期 loss 降不下去margin 过小时模型收敛快但向量空间里正负样本分界线模糊评估时 hit10 往往偏低。学习率这里有个经验判断用 Adam 时 lr 超过 0.01训练初期 loss 会上下乱跳超过 0.05 基本等于没训练loss 被 margin 卡住。我见过很多复现翻车都是因为 lr 随手填了 0.1直接把 embedding 推出合理范围。先把 lr 固定在 0.001把 margin 调到 0.5 到 2.0 之间扫一遍比同时调两个参数更容易定位问题。评估指标这里先提一句后文会专门展开FB15k 的标准评估是 link prediction用 mean rank 和 hit10 两个指标。mean rank 越小越好代表真实尾实体的排名平均越靠前hit10 是真实实体排进前 10 的比例越大越好。指标算不对前面的训练效果再好也没说服力。5. TransE 常见问题排查负样本泄露、测试集误用和随机种子玄学5.1 训练指标虚高测试指标却崩了负样本泄露现象训练 loss 降得很快甚至到 0.01 以下valid 和 test 上的 mean rank 却不降甚至比随机基线还差。原因负采样时没有过滤真实三元组或者过滤集合只用了当前 batch导致负样本里混入了 train.txt 中本来就存在的三元组。模型学到的是“别把那些假负例推得太远”真正需要压低的分数没有被压低。解决负采样前确定唯一的train_set全集候选三元组一旦命中就重新采样过滤集合用 set 而不是 listlist 的 in 判断是 O(n)会在几万样本规模下把训练拖成蜗牛。5.2 把 valid.txt 和 test.txt 混进训练数据现象valid 和 test 指标都高得吓人但换一个同类数据集跑就原形毕露。原因有些人下载压缩包后嫌三个文件麻烦直接cat train.txt valid.txt test.txt all.txt当成训练集然后又在同样的 test.txt 上评估。这是典型的“开卷考试”测试三元组在训练时已经被模型见过指标没有任何泛化意义。解决train、valid、test 三个集合在训练前做严格交集检查先确认len(set(train) set(test)) 0再开始训练。FB15k 的设计本来就是互斥划分压缩包给你的三个 txt 别合并。5.3 随机种子不固定同一份代码两次跑出不同结果现象两次训练的 mean rank 差 10 以上你以为改坏了代码其实只是随机性在作怪。原因负采样用 Python random 打乱三元组顺序Embedding 初始化用 PyTorch 的随机数两个随机源独立工作只要有一个不固定结果就不稳定。解决在脚本入口固定三个随机源import random import numpy as np import torch random.seed(42) np.random.seed(42) torch.manual_seed(42) if torch.cuda.is_available(): torch.cuda.manual_seed_all(42)random 模块负责 train_ids 的 shufflenumpy 负责可能在数据预处理里用到的随机数组torch 负责初始化缺一个都会让复现失败。这个行为不是玄学而是随机数状态树的叠加效果。做对比实验时如果不固定种子margin 和维度的差异可能被随机噪声淹没。5.4 重启进程后 entity2id 对不上向量白训一场现象训练结束保存了模型权重第二天重新加载做评估把 test.txt 转回 id 时发现实体名报 KeyError或者明明保存了 100 个维度加载的权重却和实体顺序错位。原因只保存了权重没保存 entity2id 和 relation2id 这两个映射文件。PyTorch 的 state_dict 只按向量在 embedding 表里的顺序存不包含实体名字一旦映射顺序变了权重就等于随机数。解决训练结束时把映射写进 JSON 或文本文件import json with open(entity2id.json, w, encodingutf-8) as f: json.dump(entity2id, f, ensure_asciiFalse) with open(relation2id.json, w, encodingutf-8) as f: json.dump(relation2id, f, ensure_asciiFalse) torch.save(model.state_dict(), transe_fb15k.pt)加载时先读 JSON再恢复 entity2id最后再读权重。JSON 的 key 是字符串注意把 id 转回 int否则拿字符串去查 embedding 表会报类型错误。养成“权重和映射一起存”的习惯后训练中断也能在任意 epoch 续上不至于白跑几小时。5.5 CPU 训练越来越慢负采样和评估是真正瓶颈现象epoch 1 到 10 速度正常越往后越慢或者 GPU 显存占用不高但 GPU 利用率只有 20%。原因负采样在一个大名单上反复random.choice评估阶段又用一个 Python for 循环对每个实体逐条计算 L2 距离这些 Python 层操作把 PyTorch 的加速全部吞掉了。解决负采样尽量用 NumPy 向量化生成候选索引不要把random.choice放在 PyTorch 的 autograd 图里评估阶段把实体向量整理成一个[num_entities, dim]矩阵用广播一次算出所有距离再用 argsort 取排名。FB15k 实体数一万五、关系数一千多的规模评估全部三元组确实要跑一阵子但至少应该从“一个实体一个循环”优化到“一个 batch 一个矩阵运算”。6. 用 t-SNE 验证实体向量可视化和类比推理的最后一公里训练结束后别急着交差。mean rank 和 hit10 是数值指标但嵌入向量是不是真的学到了语义结构还得靠两件事交叉验证把实体向量压到二维画出来以及用关系平移做类比推理。先用 t-SNE 看看实体向量的簇结构。FB15k 的实体名大多是 URI 风格直接看图看不出语义但同一类实体是否聚在一起能明显反映训练质量from sklearn.manifold import TSNE import matplotlib.pyplot as plt entity_matrix model.entity_emb.weight.detach().cpu().numpy() sample_ids random.sample(range(len(entity2id)), 500) sample_vecs entity_matrix[sample_ids] tsne TSNE(n_components2, perplexity30, random_state42) coords tsne.fit_transform(sample_vecs) plt.scatter(coords[:, 0], coords[:, 1], s3) plt.savefig(transe_entity_tsne.png, dpi200)t-SNE 的 perplexity 取值会影响簇的视觉效果一般设在 20 到 50 之间。这里只需要确认簇不是一团乱麻不用过度解读坐标。如果所有点糊成均匀一团多半是负采样没起作用模型根本没学出区分度。也可以用向量平移做一个简单的类比自检取一个测试三元组 (h, r, t)计算 t_emb - h_emb 的向量方向和 r_emb 做余弦相似度如果 TransE 真的学到了“头实体沿关系方向平移到达尾实体”这两个向量方向应该基本一致。这个检查比 mean rank 更直观它直接回扣 h r ≈ t 这条假设。我在做嵌入项目时吃过一个亏只盯着 hit10 调参直到画向量图才发现模型其实只学到局部共现没学到关系平移。后来我会把 t-SNE 和关系一致性检查写进训练脚本每个 epoch 结束跑一次抽样验证指标再好看也不如向量图直观。另一个让我记住教训的习惯是保存模型时连同 entity2id、relation2id 和训练用的超参一起存成 JSON恢复权重时三件套缺一不可。第一次跑通时我把 entity2id 随手存在内存里第二天换进程加载权重后彻底对不上号几小时训练作废。从那以后任何嵌入模型的项目我都先把映射文件写好再开始调参数。希望这些步骤和你自己跑出来的曲线能对上也希望这次复现过程帮你省掉几个晚上的无头排查时间。希望帮到你。本文还有配套的精品资源点击获取