Python实现TransE在FB15k上的训练：避开三大坑与调参实战

简介：这是一份TransE模型的Python实现，配套FB15k知识图谱数据集，面向知识图谱表示学习入门者、算法工程师及需要完成链接预测、知识图谱补全等任务的研究人员。资源共21个文件，包含14个txt数据文件（例如训练集、验证集与测试集的划分）、4个py源码文件、2个md说明文档以及1个png示意图，压缩包整体仅5.85MB。已有401人学习下载。源码完整实现了TransE的核心训练流程：随机初始化实体与关系向量，通过负采样构造错误三元组，采用L1或L2距离作为损失函数，利用梯度下降进行优化，使正确三元组满足头实体向量加关系向量约等于尾实体向量这一约束。基于FB15k真实数据，可直接运行训练、验证与测试，并输出MRR、HITS@k等常用评估指标。配套的README与示意图能辅助理解模型结构、数据格式和训练参数，适合作为课程设计、论文复现或实际知识图谱项目的基线，也便于后续调整向量维度、学习率、批大小等超参数开展对比实验。

1. 用 Python 把 TransE 跑在 FB15k 上，先绕开这 3 个坑

TransE 是知识图谱嵌入里最朴素也最常被当作基线模型的算法，FB15k 则是 Freebase 的一个 15k 子集，train.txt里存的就是它的标准三元组训练数据。标题里的train.txt分_TransE.zip拆开看，其实已经点明了完整工作流：把 FB15k 的三元组数据按实体、关系、训练/验证/测试切分好，再用 Python 实现 TransE 在train.txt上完成训练和评估。一个反直觉的结论是：TransE 在 FB15k 上收敛很快，但指标涨得很慢。多数人第一次跑出来的 Hits@10 和公开基线差一截，问题往往不在模型本身，而在数据划分方式、负采样策略和归一化时机这三个细节上。这篇文章面向刚接触知识图谱嵌入的工程师，也适合已经跑通代码但调不动指标的人。你会看到理论怎么落到数据上，以及哪些参数值得调、哪些坑必须绕开。

2. TransE 的训练目标与 FB15k 数据格式，先搞清楚 train.txt 里每一行是什么

2.1 TransE 的得分函数与 margin loss：h + r ≈ t 怎么变成可优化目标

TransE 的核心假设只有一句话：在一个嵌入空间里，头实体向量加上关系向量，应该约等于尾实体向量，也就是h + r ≈ t。这个假设简单到近乎粗暴，但它直接把知识图谱里的符号三元组变成了向量空间里的平移关系。得分函数最常见的写法是：

• L1 距离：||h + r - t||_1
• L2 距离：||h + r - t||_2

得分越低，说明这个三元组越合理。训练目标不是让正样本得分变成 0，而是让正样本得分比负样本低出一个 margin。这就是 margin-based ranking loss，形式是：

loss = max(0, margin + d(h + r, t) - d(h' + r, t'))

其中(h, r, t)是train.txt里的正样本三元组，(h', r, t')是把头实体或尾实体随机替换掉后生成的负样本。这个 loss 的意思是：正样本的距离要比负样本的距离小至少margin。如果差距已经够大，loss 就是 0，梯度也归零，所以 TransE 的训练过程天然是稀疏更新的。

选择 margin 时要留意一个权衡：margin 太小，正负样本的区分度不够，嵌入向量会挤成一团；margin 太大，模型会过度关注那些难分的负样本，训练不稳定。FB15k 上常见的取值在 0.5 到 2.0 之间，但这和 embedding 维度、距离函数、学习率都耦合在一起，后面第 4 章会具体说怎么组合调参。

2.2 FB15k 的 train.txt / valid.txt / test.txt 结构与 ID 映射

FB15k 标准发布包里通常有三个文件：train.txt、valid.txt、test.txt。每个文件一行一个三元组，字段之间用 tab 分隔，格式是head\trelation\ttail。注意这里不是空格，也不是逗号。拿到train.txt分_TransE.zip之后，第一步先解压看结构：

unzip train.txt分_TransE.zip wc -l train.txt valid.txt test.txt head -5 train.txt

wc -l统计行数，能快速确认三个文件的规模比例。FB15k 常见的划分是训练集约 47 万条、验证集约 5 万条、测试集约 5.9 万条，但不同版本的预处理可能略有出入，以你实际解压出来的行数为准。head -5看前五行长什么样，确认是head\trelation\ttail而不是别的顺序。

实体和关系在原始文件里是字符串名字，比如/m/02_m8jq这种 Freebase 风格的 ID。不能直接把字符串喂给神经网络，需要先做 ID 映射。常见做法是遍历所有文件，把出现过的实体和关系各编一个从 0 开始的整数 ID：

entities, relations = set(), set() for fname in ['train.txt', 'valid.txt', 'test.txt']: with open(fname, 'r', encoding='utf-8') as f: for line in f: h, r, t = line.strip().split('\t') entities.add(h) entities.add(t) relations.add(r) entity2id = {e: i for i, e in enumerate(entities)} relation2id = {r: i for i, r in enumerate(relations)}

这段代码把所有三元组的头实体、尾实体收进entities集合，关系收进relations集合，然后分别建映射。enumerate从 0 开始编号，ID 连续且紧凑。这样做的原因是 embedding 层本质是一个查找表，ID 必须连续，否则nn.Embedding会报索引越界或者浪费大量内存。

2.3 数据划分：train.txt 分出来之后，怎么生成训练和评估要用的张量

第 2.2 节建好了 ID 映射，接下来要把train.txt转成 PyTorch 能直接读取的 tensor。这里有一个很多人会忽略的坑：验证集和测试集的负样本必须排除掉训练集里出现过的三元组。FB15k 的标准做法是 filtered 评估，也就是在计算排名时，把训练集里已有的其他正确三元组从候选里去掉，否则 Hits@10 会被严重高估。这个逻辑在数据准备阶段就要留下训练集三元组的集合，后面评估要用。

import torch def load_triples(fname, entity2id, relation2id): triples = [] with open(fname, 'r', encoding='utf-8') as f: for line in f: h, r, t = line.strip().split('\t') triples.append((entity2id[h], relation2id[r], entity2id[t])) return torch.tensor(triples, dtype=torch.long)

load_triples输出的 tensor 形状是(num_triples, 3)，每一行是(head_id, relation_id, tail_id)。训练时按 batch 切分这个 tensor，评估时再单独处理。到这一步，train.txt已经变成了模型可以消费的数字形式，后面的实现都建立在这套 ID 体系上。

3. Python 实现 TransE 的最小可训练版本，用 PyTorch 从零写训练闭环

3.1 数据加载器：把 train.txt 按 batch 喂给模型并生成负样本

数据准备阶段拿到的是完整训练三元组 tensor，训练时要随机打乱并按 batch 切分。PyTorch 的DataLoader配合TensorDataset可以做这件事，但负样本生成需要自己写。常见的负采样做法是：对 batch 里的每个正样本，随机替换头实体或尾实体，替换的实体从全量实体集合里均匀采样。下面这个NegativeSampler是我常用的写法：

class NegativeSampler: def __init__(self, num_entities): self.num_entities = num_entities def corrupt(self, heads, tails, bern_prob=0.5): heads = heads.clone() tails = tails.clone() mask = torch.rand(heads.size(0)) < bern_prob heads[mask] = torch.randint(0, self.num_entities, (mask.sum(),)) tails[~mask] = torch.randint(0, self.num_entities, ((~mask).sum(),)) return heads, tails

corrupt方法接收正样本的heads和tails向量，按bern_prob的概率决定是破坏头还是破坏尾。mask是布尔张量，True的位置替换头实体，False的位置替换尾实体。torch.randint在[0, num_entities)区间均匀采样，这个范围必须覆盖全部实体 ID，否则会出现采样不到某些实体的情况。

3.2 模型定义与参数初始化：nn.Embedding 和 L2 归一化的边界

TransE 的模型结构是一个双 embedding 表：一个给实体，一个给关系。实体 embedding 的形状是(num_entities, dim)，关系 embedding 的形状是(num_relations, dim)。PyTorch 实现如下：

import torch.nn as nn class TransE(nn.Module): def __init__(self, num_entities, num_relations, dim): super().__init__() self.entities = nn.Embedding(num_entities, dim) self.relations = nn.Embedding(num_relations, dim) nn.init.uniform_(self.entities.weight, -1.0, 1.0) nn.init.uniform_(self.relations.weight, -1.0, 1.0) def forward(self, heads, relations, tails, normalize=True): h = self.entities(heads) r = self.relations(relations) t = self.entities(tails) if normalize: h = torch.nn.functional.normalize(h, p=2, dim=1) t = torch.nn.functional.normalize(t, p=2, dim=1) return h, r, t

初始化用均匀分布[-1, 1]，这是 TransE 原论文里的做法。关键点在normalize=True：头实体和尾实体的 embedding 每步都要做 L2 归一化，但关系向量不归一化。这样做的原因是，实体的模长差异如果不被约束，模型会倾向于让某些实体的模长变得特别大，从而让距离度量失效。关系不归一化是因为关系的模长本身携带了语义信息，归一化会把这种信息抹掉。这个边界很微妙，后面第 5 章还会展开。

3.3 训练循环与损失计算：SGD + margin loss 的完整代码

模型定义好了，负样本也有生成器了，接下来是训练主循环。损失函数用第 2.1 节的 margin ranking loss，优化器一般用SGD而不是 Adam。Adam 对 TransE 效果不稳定，因为它的自适应学习率会放大稀疏梯度的影响，而 TransE 的梯度本来就是稀疏的：

def train(model, train_triples, sampler, optimizer, margin, batch_size=1024, epochs=1000): model.train() num_train = train_triples.size(0) for epoch in range(epochs): perm = torch.randperm(num_train) train_triples = train_triples[perm] total_loss = 0.0 for batch_start in range(0, num_train, batch_size): batch = train_triples[batch_start:batch_start + batch_size] heads, relations, tails = batch[:, 0], batch[:, 1], batch[:, 2] neg_heads, neg_tails = sampler.corrupt(heads, tails) h, r, t = model(heads, relations, tails) h_neg, r_neg, t_neg = model(neg_heads, relations, neg_tails) pos_dist = (h + r - t).norm(p=2, dim=1) neg_dist = (h_neg + r_neg - t_neg).norm(p=2, dim=1) loss = torch.clamp(margin + pos_dist - neg_dist, min=0).mean() optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() if epoch % 100 == 0: print(f'epoch {epoch}, loss {total_loss:.4f}')

每个参数都要说清楚含义。perm是torch.randperm生成的随机排列索引，用它重排训练三元组，保证每个 epoch 看到的 batch 组成不同。batch_size=1024是兼顾内存和梯度稳定性的常见取值。pos_dist计算正样本的 L2 距离，neg_dist计算负样本的 L2 距离，torch.clamp(..., min=0)实现max(0, margin + pos - neg)的合页损失。.mean()取 batch 内平均，让 loss 数值和 batch_size 无关。optimizer.zero_grad()是必须的，否则梯度会跨 batch 累积。

优化器的初始学习率一般取0.01到0.1之间，比常规神经网络的默认值高一个量级。原因是 TransE 的目标函数相对简单，损失曲面没有那么多局部最优，大学习率反而能更快找到一个合理的嵌入空间：

optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

4. 在 FB15k 上调参：embedding 维数、margin、batch size 怎么影响指标

4.1 三个最影响结果的参数：dim、margin、学习率

FB15k 上的表现很大程度上被三个超参数决定：embedding 维度dim、损失函数里的margin、优化器的学习率。它们之间的耦合关系比大多数人以为的更强。先看一个实际参考区间：

dim的选择要看实体和关系的数量。FB15k 实体约 1.5 万、关系约 1300 个，100 维在表达力和计算开销之间比较平衡。如果你把 dim 提到 200，训练时间接近翻倍，但 Hits@10 往往只提升一两个点，甚至可能因为过拟合而下降。margin的设置和距离函数强耦合：用 L2 距离时，距离值的量级比 L1 大，margin 通常取 1.0 以上；用 L1 距离时，margin 取 0.5 到 1.0 就够了。

4.2 用 valid.txt 做 early stopping，并算 MRR 和 Hits@10

训练 TransE 不能只盯着训练 loss。很多时候训练 loss 一路下降，但验证集指标早就开始掉头了，这就是典型的过拟合。要避免这种状况，需要在每个 epoch 结束后用valid.txt评估一次指标，并且只保存验证指标最好的那一次模型权重。评估要同时算 MRR 和 Hits@10，这两个指标在 FB15k 上是标准配置：

def evaluate(model, valid_triples, all_triples, num_entities): model.eval() ranks = [] with torch.no_grad(): for h, r, t in valid_triples: h, r, t = h.item(), r.item(), t.item() h_emb = model.entities(torch.tensor([h])) r_emb = model.relations(torch.tensor([r])) t_emb = model.entities.weight scores = (h_emb + r_emb - t_emb).norm(p=2, dim=1) # filtered: 排除训练集中存在的其他正确三元组 for e in range(num_entities): if e == t: continue if (h, r, e) in all_triples: scores[e] = float('inf') rank = (scores < scores[t]).sum().item() + 1 ranks.append(rank) mrr = (1.0 / torch.tensor(ranks)).mean().item() hits10 = sum(1 for r in ranks if r <= 10) / len(ranks) return mrr, hits10

这里我在评估时把尾实体全部替换为候选实体，计算每个候选的得分，然后找到正确尾实体得分排第几。filtered 的部分很关键：if (h, r, e) in all_triples时把这个候选的分数设成inf，意思是跳过那些在训练集里真实存在但不是当前尾实体的三元组。如果不做这一步，MRR 会虚高。all_triples要在开始训练之前就把train.txt的全部三元组转成集合传进来。

4.3 FB15k 上的常见坑：训练 loss 下降但验证指标不涨怎么排查

症状一：训练 loss 下降但 MRR 和 Hits@10 完全不动。先看负采样是不是有问题。如果corrupt时替换了头实体但新头实体恰好和原头实体是同一个实体，这个负样本就是假的。虽然概率只有1/num_entities，但 FB15k 实体数不多，假负样本会把模型往错误方向推。症状二：验证集 Hits@10 和训练集 Hits@10 差距很大，训练集上接近 0.9、验证集只有 0.3，这是过拟合。此时优先调小dim或提前停止训练，而不是加数据。症状三：训练 loss 一开始就震荡，很可能是学习率太大，把学习率从 0.1 降到 0.02 再试。

还有一个在 FB15k 上特别容易踩的坑：实体 embedding 归一化之后，embedding 向量全部落在单位球面上，距离的数值范围被压得很小。这时如果 margin 设置得比最大可能距离还大，loss 会恒大于零，模型会一直尝试拉开距离但永远做不到。排查方法是打印几轮pos_dist和neg_dist的均值，如果正样本距离已经在 0.5 以下但 margin 是 2.0，说明 margin 设大了。

5. 在 train.txt 上把数据价值榨干：bern 负采样、归一化时机与两阶段训练

5.1 为什么均匀负采样在 FB15k 上不够，改成 bern 采样

第 3.1 节的均匀负采样在数学上是无偏的，但对 FB15k 这种真实知识图谱来说不够高效。原因在于不同关系的头实体和尾实体分布差异极大。举一个具体例子：/location/location/contains这种关系，头实体通常是国家或城市，尾实体通常是景点或区域，实体类型差异明显。如果用均匀采样替换头实体，很可能采到一个本质上不可能作为头实体的实体，这种负样本太简单，模型学不到东西。

Bernoulli 负采样策略会统计每个关系下头实体被替换和尾实体被替换的频率，然后按这个频率分配替换概率。具体做法是：训练前统计每个关系r下，正样本三元组(h, r, t)中每个头实体h平均对应多少个尾实体，以及每个尾实体t平均对应多少个头实体，然后算出两个替换概率。实现上要额外维护两个计数表：

def count_bern_probs(train_triples, num_relations): head_count = torch.zeros(num_relations) tail_count = torch.zeros(num_relations) rel_count = torch.zeros(num_relations) for h, r, t in train_triples: head_count[r] += 1 tail_count[r] += 1 rel_count[r] += 1 head_prob = head_count / (head_count + tail_count) return head_prob

把head_prob传到负采样器里，替换头实体的概率就不再是固定的 0.5，而是由数据分布决定的动态值。FB15k 上用 bern 采样通常能让 Hits@10 提高 2 到 5 个点，提升幅度甚至比调 embedding 维度还明显。

5.2 L2 归一化的时机：训练完成后再归一化不会提升指标

实体 embedding 的 L2 归一化不止是一个「要不要做」的问题，更关键的是「什么时候做」。常见做法是在 forward 里对实体的 embedding 做归一化，也就是第 3.2 节代码里的normalize=True。如果改成训练时不归一化、训练完再统一归一化，指标不会变好。原因在于距离的数值分布已经被训练过程塑形了，事后归一化只是把所有向量等比例缩放到单位球面，相对距离排序一点不变。

另外需要明确：只有实体需要归一化，关系永远不归一化。验证方法是训练完成后打印实体 embedding 和关系 embedding 的模长分布，你会看到实体模长全为 1，而关系模长分布很广。如果关系也被归一化了，得分函数的表达能力会被限制在一个球面上，无法表达模长差异带来的语义丰富性。

5.3 一个实用技巧：先粗训再精调，用早停的模型做冷启动

FB15k 上我一般会跑两阶段训练。第一阶段用较粗的参数快速找到大致合理的区域：dim=50, margin=1.0, lr=0.1, batch_size=2048，训练到验证集 Hits@10 不再上升就停，通常 200 到 400 个 epoch 就够。这个阶段不求指标高，只求找到一个不会太差的初始点。第二阶段做精调：把 embedding 维度提高到目标值，初始化用第一阶段学到的向量插值到高维空间，学习率降到0.01，margin 改成1.5，继续训练。

两阶段训练的好处是，第一阶段能快速鉴别数据预处理和代码有没有 bug。如果阶段一在几十个 epoch 内 Hits@10 一直小于 0.1，别急着调参，先回头看负采样有没有问题，看 ID 映射有没有错位。等模型能稳定跑到 0.2 以上，再进入第二阶段精调。这时候每个 epoch 的训练速度已经快了很多，因为权重的初始化已经处于一个较优区域，梯度更新的步长可以更小，收敛也更平稳。最后保存验证集 MRR 最高的 checkpoint，用同一份权重在test.txt上跑一次最终指标，作为对外报告的结果。检查test.txt里的三元组是否有和训练集完全重复的行，如果有，评估前直接去掉，否则 Hits@10 会被异常拉高。

本文还有配套的精品资源，点击获取