性能技巧

PyKEEN 使用一系列技术组合来促进训练/评估期间的高效计算，并尝试最大化可用硬件的利用率（目前主要关注单GPU的使用）。

实体和关系ID

三元组中的实体和关系通常以字符串形式存储。因为KGEMs的目标是学习这些实体和关系的向量表示，使得所选的交互函数能够在它们之上学习到有用的评分，我们需要一个从字符串表示到向量的映射。此外，为了计算效率，我们希望将所有实体/关系嵌入存储在矩阵中。因此，映射过程包括两部分：将字符串映射到ID，并使用ID访问嵌入（=行索引）。

在PyKEEN中，映射过程发生在pykeen.triples.TriplesFactory。三元组工厂维护唯一的实体和关系标签集，并确保它们映射到唯一的整数ID上 \([0,\text{num_unique_entities})\) 对于实体和 \([0, \text{num_unique_relations})\)。映射分别可通过属性 :data:pykeen.triples.TriplesFactory.entity_label_to_id 和 :data:pykeen.triples.TriplesFactory.relation_label_to_id 访问。

为了提高性能，映射过程只发生一次，并且基于ID的三元组存储在张量中：pykeen.triples.TriplesFactory.mapped_triples。

元组广播

交互函数通常仅用于评分单个三元组 \((h, r, t)\) 的标准情况。此函数在 PyKEEN 中通过每个模型的 pykeen.models.base.Model.score_hrt() 方法实现，例如 pykeen.models.DistMult 的 pykeen.models.DistMult.score_hrt()。在局部封闭世界假设（LCWA）下进行训练、评估模型以及执行链接预测任务时，目标是为给定元组（即 \((h, r)\)、\((r, t)\) 或 \((h, t)\)）的所有实体/关系进行评分。在这些情况下，单个元组会多次用于不同的实体/关系。

例如，我们想要使用pykeen.models.DistMult对单个元组\((h, r)\)的所有实体进行排名，该元组来自pykeen.datasets.FB15k237。这个数据集包含14,505个实体，这意味着有14,505个\((h, r, t)\)组合，其中\(h\)和\(r\)是常数。查看pykeen.models.DistMult的交互函数，我们可以观察到\(h \odot r\)部分导致计算\(h \odot r \odot t\)的数学操作减少了一半。因此，只计算一次\(h \odot r\)部分并重复使用它，可以为我们节省其他14,504个实体的数学操作，使计算总体上大约快两倍。在广播部分相对于整个交互函数具有较高相对复杂性的情况下，速度提升可能会显著更高，例如pykeen.models.ConvE。

为了使这种技术成为可能，PyKEEN模型必须通过模型类中的以下方法提供一个显式的广播函数：

pykeen.models.base.Model.score_h() - 为给定的\((r, t)\)元组评分所有可能的头实体

pykeen.models.base.Model.score_r() - 为给定的\((h, t)\)元组评分所有可能的关系

pykeen.models.base.Model.score_t() - 为给定的\((h, r)\)元组评分所有可能的尾实体

PyKEEN架构原生支持这些方法，并在可能的情况下使用这种技术，无需任何额外的修改。提供这些方法是完全可选的，在实现新模型时不需要。

基于索引掩码的过滤

在这个例子中，给定了一个知识图谱 \(\mathcal{K} \subseteq \mathcal{E} \times \mathcal{R} \times \mathcal{E}\) 以及训练三元组 \(\mathcal{K}_{train}\)、测试三元组 \(\mathcal{K}_{test}\) 和验证三元组 \(\mathcal{K}_{val}\) 中的 \(\mathcal{K}\) 的不相交并集。对 \(\mathcal{K}_{test}\) 和 \(\mathcal{K}_{val}\) 执行相同的操作，但本节仅以 \(\mathcal{K}_{test}\) 为例。

在知识图谱嵌入模型的标准评估过程中，对于每个测试三元组 \((h, r, t) \in \mathcal{K}_{test}\)，使用交互函数 \(f:\mathcal{E} \times \mathcal{R} \times \mathcal{E} \rightarrow \mathbb{R}\) 进行链接预测任务时，会执行两个计算：

\((h, r)\) 与所有可能的尾实体 \(t' \in \mathcal{E}\) 结合，形成三元组 \(T_{h,r} = \{(h,r,t') \mid t' \in \mathcal{E}\}\)
\((r, t)\) 与所有可能的头实体 \(h' \in \mathcal{E}\) 结合形成三元组 \(H_{r,t} = \{(h',r,t) \mid h' \in \mathcal{E}\}\)

最后，\((h, r, t)\) 的排名是针对所有 \((h, r, t') \in T_{h,r}\) 和 \((h', r, t) \in H_{r,t}\) 三元组相对于交互函数 \(f\) 计算的。

在过滤设置中，\(T_{h,r}\) 不允许包含尾实体 \((h, r, t') \in \mathcal{K}_{train}\)，并且 \(H_{r,t}\) 不允许包含导致 \((h', r, t) \in \mathcal{K}_{train}\) 三元组的头实体在训练数据集中找到的。因此，它们的定义可以修改为：

\(T^{\text{filtered}}_{h,r} = \{(h,r,t') \mid t' \in \mathcal{E}\} \setminus \mathcal{K}_{train}\)
\(H^{\text{filtered}}_{r,t} = \{(h',r,t) \mid h' \in \mathcal{E}\} \setminus \mathcal{K}_{train}\)

虽然这在理论上很容易定义，但它带来了几个实际挑战。例如，它导致了计算上的挑战，即所有新的可能三元组 \((h, r, t') \in T_{h,r}\) 和 \((h', r, t) \in H_{r,t}\) 必须被枚举，然后在 \(\mathcal{K}_{train}\) 中检查是否存在。考虑到像 pykeen.datasets.FB15k237 这样的数据集，它有近15,000个实体，每个测试三元组 \((h,r,t) \in \mathcal{K}_{test}\) 会导致 \(2 * | \mathcal{E} | = 30,000\) 个可能的新三元组，这些三元组必须在训练数据集中进行检查，然后被移除。

为了获得非常快速的过滤，PyKEEN 结合了上述在实体和关系ID 和元组广播中介绍的技术，以及以下机制，这在我们的情况下使得过滤评估的速度比之前版本中使用的机制提高了600,000倍。

作为起点，PyKEEN 将始终计算 \(H_{r,t}\) 和 \(T_{h,r}\) 中所有三元组的分数，即使在过滤设置中也是如此。因为平均而言，正三元组的数量非常低，所以需要移除的结果很少。此外，由于在 Tuple Broadcasting 中提出的技术，对额外实体进行评分的成本非常低。因此，我们从 pykeen.models.base.Model.score_t() 对所有三元组 \((h, r, t') \in H_{r,t}\) 和 pykeen.models.base.Model.score_h() 对所有三元组 \((h', r, t) \in T_{h,r}\) 的分数向量开始。

接下来，创建稀疏过滤器 \(\mathbf{f}_t \in \mathbb{B}^{| \mathcal{E}|}\) 和 \(\mathbf{f}_h \in \mathbb{B}^{| \mathcal{E}|}\)，它们表示哪些实体会导致在训练数据集中找到的三元组。为了实现这一点，我们将依赖于在实体和关系ID 中介绍的技术，即所有实体/关系ID对应于它们在各自嵌入张量中的确切位置。例如，我们从测试三元组 \((h, r, t) \in \mathcal{K}_{test}\) 中取出元组 \((h, r)\)，并关注所有应从 \(T_{h,r}\) 中移除的尾实体 \(t'\)，以获得 \(T^{\text{filtered}}_{h,r}\)。这是通过执行以下步骤实现的：

取\(r\)并将其与训练数据集中所有三元组的关系进行比较，生成一个布尔向量，其大小为训练数据集中包含的三元组数量，当任何三元组具有关系\(r\)时为真。
取\(h\)并将其与训练数据集中所有三元组的头实体进行比较，生成一个布尔向量，其大小为训练数据集中包含的三元组数量，当任何三元组的头实体为\(h\)时，该位置为真。
结合两个布尔向量，生成一个布尔向量，其大小为训练数据集中包含的三元组数量，当任何三元组同时具有头实体\(h\)和关系\(r\)时为真。
将布尔向量转换为非零索引向量，指出训练数据集在哪些索引处包含同时包含头实体h和关系\(r\)的三元组，其大小为非零元素的数量
索引向量现在应用于训练数据集的尾部实体列，返回所有与\(h\)和\(r\)结合导致三元组包含在训练数据集中的尾部实体ID \(t'\)
最后，\(t'\) 尾部实体ID索引向量应用于最初由 pykeen.models.base.Model.score_t() 返回的向量，针对所有可能的三元组 \((h, r, t')\)，并且所有受影响的分数根据IEEE-754规范设置为 float('nan')，这使得这些分数不可比较，有效地导致所有可能的新三元组 \((h, r, t') \in T^{\text{filtered}}_{h,r}\) 的分数向量。

\(H^{\text{filtered}}_{r,t}\) 是从 \(H_{r,t}\) 以类似的方式获得的。

子批次处理与切片

随着模型和数据集的规模不断增长，手头的KGEM可能会超过GPU提供的内存。特别是在训练过程中，可能希望使用一定的批量大小进行训练。当这个批量大小对于手头的硬件来说太大时，PyKEEN允许设置一个在\([1, \text{batch_size}]\)范围内的子批量大小。当设置了子批量大小时，PyKEEN会在每个子批量后自动累积梯度，并在训练期间清除计算图。这使得可以在GPU上训练那些对于手头硬件来说太大的KGEM，同时获得的结果与不使用子批量训练的结果相同。

注意

为了保证等效的结果，并非所有模型都支持子批次处理，因为某些组件，例如批量归一化，需要一次性计算整个批次以避免改变统计信息。

注意

子批次处理有时也被称为梯度累积，例如，由huggingface的 transformer 库所使用，因为我们在更新参数之前会在多个子批次上累积梯度。

对于一些大型配置，即使应用了子批次技巧，仍然可能发生内存不足的错误。在这种情况下，PyKEEN 实现了另一种技术，称为切片。请注意，我们通常为每个批次元素计算多个分数：在 sLCWA 中，我们有 \(1 + \text{num_negative_samples}\) 个分数，而在 LCWA 中，我们为每个批次元素有 \(\text{num_entities}\) 个分数。在切片中，我们不会一次性计算所有这些分数，而是在较小的“批次”中计算。对于旧式模型，即那些从 pykeen.models.base._OldAbstractModel 继承的模型，必须为每个模型单独实现这一点。新式模型，即那些从 pykeen.models.nbase.ERModel 派生的模型，有一个通用的实现，可以为所有交互启用切片。

注意

切片计算以较小的批次计算分数，但仍需要计算所有分数的梯度，因为一些损失函数需要访问它们。

自动内存优化

在训练和评估过程中，确保不超过可用硬件内存的同时允许高计算吞吐量，需要了解当前模型配置下可能的最大训练和评估批量大小。然而，确定训练和评估批量大小是一个繁琐的过程，并且在运行大量异构实验时不可行。因此，PyKEEN 在开始实际计算之前，有一个自动内存优化步骤，计算当前模型配置和可用硬件下可能的最大训练和评估批量大小。如果用户提供的批量大小对于使用的硬件来说太大，自动内存优化会确定训练的最大子批量大小，并通过上述过程累积梯度子批量 & 切片。批量大小是通过二分搜索确定的，考虑到 CUDA 架构，确保所选的批量大小是最具 CUDA 效率的。

通常评估是在GPU上进行的，以获得更快的速度。感谢torch_max_mem，你可以完全将寻找最大批量大小的任务交给PyKEEN。此外，用户可能会在评估配置中预先选择一个批量大小，以充分利用GPU，实现尽可能快的评估速度。然而，在较大的设置中测试不同的模型配置和数据集分区（例如HPO）时，硬件需求可能会急剧变化，这可能导致评估无法再使用预设的批量大小运行，或者对于较大的数据集和内存密集型模型，根本无法在GPU上运行。在这些情况下，torch_max_mem将负责降低实际批量大小，直到不再发生内存不足的错误。

评估回退

在某些情况下，即使使用最小的批量大小和切片，评估也可能无法在GPU上成功。在这些罕见的情况下，PyKEEN提供了回退到CPU的选项。注意：在评估回退到使用CPU的情况下，这可能会导致评估时间显著延长。