NodePiece入门指南

本页面提供了更多关于使用和配置NodePiece的实际示例。

基本用法

我们将使用pykeen.datasets.FB15k237用于说明目的 在以下示例中。

from pykeen.models import NodePiece
from pykeen.datasets import FB15k237

# inverses are necessary for the current version of NodePiece
dataset = FB15k237(create_inverse_triples=True)

在最简单的pykeen.models.NodePiece使用中，我们只会使用关系进行标记化。我们可以通过以下参数来实现这一点：

  1. 将tokenizers="RelationTokenizer"设置为 pykeen.nn.node_piece.RelationTokenizer。我们可以简单地引用 类名，它会自动解析为pykeen.nn.node_piece.Tokenizer的正确子类 通过class_resolver。

  2. 设置num_tokens=12以对每个节点采样12个唯一关系。如果某些实体的唯一关系少于12个，则差异将用辅助填充标记进行填充。

代码看起来是这样的：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers="RelationTokenizer",
    num_tokens=12,
    embedding_dim=64,
)

接下来，我们将使用一组分词器 (pykeen.nn.node_piece.AnchorTokenizer 和 pykeen.nn.node_piece.RelationTokenizer) 来复制完整的 NodePiece 分词，使用 \(k\) 个锚点和 \(m\) 个关系上下文。这就像 将分词器列表发送到 tokenizers 并将参数列表发送到 num_tokens 一样简单：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    embedding_dim=64,
)

类解析器将自动实例化 pykeen.nn.node_piece.AnchorTokenizer，每个节点有20个锚点， 以及pykeen.nn.node_piece.RelationTokenizer，每个节点有12个关系， 因此指定tokenizers和num_tokens的顺序在这里很重要。

锚点选择与搜索

pykeen.nn.node_piece.AnchorTokenizer 有两个字段：

  1. selection 控制我们如何从图中采样锚点（默认32个锚点）

  2. searcher 控制我们如何使用选定的锚点对节点进行标记 （默认情况下使用 pykeen.nn.node_piece.CSGraphAnchorSearcher）

默认情况下，我们上面的模型使用32个锚点，这些锚点被选为具有最高度数的节点，使用pykeen.nn.node_piece.DegreeAnchorSelection（这些是锚点选择解析器的默认值），并且节点使用pykeen.nn.node_piece.CSGraphAnchorSearcher进行标记化 - 它使用scipy.sparse以确定性的方式显式计算图中所有节点到所有锚点的最短路径。我们可以承受FB15k237大小的相对较小的图。

对于较大的图，我们建议使用广度优先搜索（BFS）程序在pykeen.nn.node_piece.ScipySparseAnchorSearcher中 - 它通过迭代扩展节点邻域来应用BFS，直到找到所需数量的锚点 - 这在像pykeen.datasets.OGBWikiKG2这样的大规模图上显著节省了计算时间。

对于拥有15k个节点的FB15k237来说，32个独特的锚点可能有点太小了 - 所以让我们创建一个pykeen.models.NodePiece模型，通过发送tokenizers_kwargs列表来选择100个使用最高度策略的锚点：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    tokenizers_kwargs=[
        dict(
            selection="Degree",
            selection_kwargs=dict(
                num_anchors=100,
            ),
            searcher="CSGraph",
        ),
        dict(),  # empty dict for the RelationTokenizer - it doesn't need any kwargs
    ],
    embedding_dim=64,
)

tokenizers_kwargs 期望的字典数量与您使用的分词器数量相同，因此我们这里有2个字典 - 一个用于 AnchorTokenizer，另一个用于 RelationTokenizer（但这个不需要任何kwargs，所以我们只放了一个空字典）。

让我们创建一个模型，使用BFS策略选择500个顶级页面排名锚点 - 我们只需修改selection和searcher参数：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    tokenizers_kwargs=[
        dict(
            selection="PageRank",
            selection_kwargs=dict(
                num_anchors=500,
            ),
            searcher="ScipySparse",
        ),
        dict(),  # empty dict for the RelationTokenizer - it doesn't need any kwargs
    ],
    embedding_dim=64,
)

看起来不错，但请系好安全带 🚀 - 我们可以依次使用几种锚点选择策略来选择更多样化的锚点！令人惊叹 😍

让我们创建一个包含500个锚点的模型，其中50%将是高度节点，另外50%将是高PageRank节点 - 为此我们有一个pykeen.nn.node_piece.MixtureAnchorSelection类！

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    tokenizers_kwargs=[
        dict(
            selection="MixtureAnchorSelection",
            selection_kwargs=dict(
                selections=["degree", "pagerank"],
                ratios=[0.5, 0.5],
                num_anchors=500,
            ),
            searcher="ScipySparse",
        ),
        dict(),  # empty dict for the RelationTokenizer - it doesn't need any kwargs
    ],
    embedding_dim=64,
)

现在selection_kwargs控制我们将使用哪些策略以及每个策略将采样多少个锚点 - 在我们的例子中是selections=['degree', 'pagerank']。使用ratios参数，我们控制这些采样锚点在总池中的比例 - 在我们的例子中是ratios=[0.5, 0.5]，这意味着degree和pagerank策略各自将从总锚点数中采样50%。由于总数为500，将有250个最高度数的锚点和250个最高pagerank的锚点。ratios 必须总和为1.0

重要: 采样的锚点是唯一的 - 也就是说，如果一个节点出现在前K度数和前K页面排名中，它只会被使用一次，采样器在后续策略中会跳过它。

目前，我们有3种锚点选择策略：degree、 pagerank和random。后者只是随机选择节点作为锚点。

让我们创建一个在原始NodePiece论文中报告的FB15k237的分词设置，其中包含40%的顶级度锚点，40%的顶级页面排名，以及20%的随机锚点：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    tokenizers_kwargs=[
        dict(
            selection="MixtureAnchorSelection",
            selection_kwargs=dict(
                selections=["degree", "pagerank", "random"],
                ratios=[0.4, 0.4, 0.2],
                num_anchors=500,
            ),
            searcher="ScipySparse",
        ),
        dict(),  # empty dict for the RelationTokenizer - it doesn't need any kwargs
    ],
    embedding_dim=64,
)

关于锚点距离的说明：目前，锚点距离是隐式考虑的，即在通过最短路径或BFS执行实际分词时，我们会根据接近程度对锚点进行排序并保留前K个最近的锚点。尚未实现将锚点距离嵌入作为位置特征添加到锚点嵌入中。

我的图表需要多少个总锚点 num_anchors 和锚点与关系 num_tokens？

这是一个具有深刻理论意义的好问题，涉及到像k-支配集和顶点覆盖集这样的NP难问题。我们没有针对每个可能的数据集的闭式解，但我们发现了一些经验启发式方法：

  • 将 num_anchors 保持为图中总节点数的1-10%是一个良好的开始

  • 图的密度是一个主要因素：图越密集，需要的num_anchors就越少。对于密集的FB15k237，总共100个锚点（超过15k个节点）似乎已经足够，而对于较稀疏的WN18RR，我们至少需要500个锚点（超过40k个节点）。对于密集的OGB WikiKG2，2.5M个节点中20K个锚点的词汇（小于1%）已经能够达到SOTA结果。

  • 同样适用于每个节点的锚点：对于更稀疏的图，您需要更多的令牌，而对于更密集的图，则需要更少。

  • 关系上下文的大小取决于图中唯一关系的密度和数量，例如，在FB15k237中我们有237 * 2 = 474个唯一关系，而在WN18RR中只有11 * 2 = 22个。如果我们选择一个太大的上下文，大多数标记将是PADDING_TOKEN，我们不希望这样。

  • 在NodePiece论文中报告的关系上下文大小（每个节点的关系数）是第66百分位数的每个节点的唯一入射关系数，例如FB15k237为12，WN18RR为5

在某些任务中，您可能根本不需要锚点，只需使用RelationTokenizer即可！查看论文以获取更多结果。

  • 在归纳链接预测任务中，我们不使用锚点，因为推理图与训练图是断开的；

  • 在关系预测中，我们发现仅使用关系上下文比使用锚点加关系更好；

  • 在节点分类（目前，此管道在PyKEEN中不可用）中，在像Wikidata这样的密集关系丰富的图上，我们发现仅关系上下文比锚点+关系更好。

使用NodePiece与pykeen.pipeline.pipeline()

让我们将最后一个NodePiece模型打包到管道中：

import torch.nn

from pykeen.models import NodePiece
from pykeen.pipeline import pipeline

result = pipeline(
    dataset="fb15k237",
    dataset_kwargs=dict(
        create_inverse_triples=True,
    ),
    model=NodePiece,
    model_kwargs=dict(
        tokenizers=["AnchorTokenizer", "RelationTokenizer"],
        num_tokens=[20, 12],
        tokenizers_kwargs=[
            dict(
                selection="MixtureAnchorSelection",
                selection_kwargs=dict(
                    selections=["degree", "pagerank", "random"],
                    ratios=[0.4, 0.4, 0.2],
                    num_anchors=500,
                ),
                searcher="ScipySparse",
            ),
            dict(),  # empty dict for the RelationTokenizer - it doesn't need any kwargs
        ],
        embedding_dim=64,
        interaction="rotate",
    ),
)

预计算词汇表

我们有一个pykeen.nn.node_piece.PrecomputedPoolTokenizer，它可以使用从本地文件或可下载链接预计算的词汇表进行实例化。

对于本地文件，指定 path：

precomputed_tokenizer = tokenizer_resolver.make(
    "precomputedpool", path=Path("path/to/vocab.pkl")
)

model = NodePiece(
    triples_factory=dataset.training,
    num_tokens=[20, 12],
    tokenizers=[precomputed_tokenizer, "RelationTokenizer"],
)

对于远程文件，请指定url：

precomputed_tokenizer = tokenizer_resolver.make(
    "precomputedpool", url="http://link/to/vocab.pkl"
)

通常，pykeen.nn.node_piece.PrecomputedPoolTokenizer 可以使用任何 pykeen.nn.node_piece.PrecomputedTokenizerLoader 作为词汇格式的自定义处理器。目前有一个这样的加载器，pykeen.nn.node_piece.GalkinPrecomputedTokenizerLoader，它期望以下格式的字典：

node_id: {
    "ancs": [a list of used UNMAPPED anchor nodes sorted from nearest to farthest],
    "dists": [a list of anchor distances for each anchor in ancs, ascending]
}

截至目前，我们不使用锚点距离，但我们期望ancs中的锚点已经按从近到远排序，因此预计算词汇表的示例可以是：

1: {'ancs': [3, 10, 5, 9, 220, ...]}  # anchor 3 is the nearest for node 1
2: {'ancs': [22, 37, 14, 10, ...]}  # anchors 22 is the nearest for node 2

未映射的锚点意味着锚点ID与实体总集中的节点ID相同0... N-1。在pickle处理过程中，我们将它们转换为一个连续的范围0 ... num_anchors-1。列表中的任何负索引将被视为填充标记（我们在预计算的词汇表中使用了-99）。

原始的NodePiece仓库有一个示例，展示了如何为OGB WikiKG 2构建这样的词汇表格式。

配置交互函数

你可以使用PyKEEN中几乎任何交互函数作为评分函数！默认情况下，NodePiece使用DistMult，但像在任何pykeen.models.ERModel中一样，很容易更改，让我们使用RotatE交互：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    interaction="rotate",
    embedding_dim=64,
)

对于RotatE，我们可能希望将关系初始化为相位 (init_phases)，并使用额外的关系约束器来保持 |r| = 1 (complex_normalize)，并使用xavier_uniform_进行 锚点嵌入初始化 - 我们也添加这个：

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    embedding_dim=64,
    interaction="rotate",
    relation_initializer="init_phases",
    relation_constrainer="complex_normalize",
    entity_initializer="xavier_uniform_",
)

配置聚合函数

本节介绍aggregation关键字参数。这是一个编码器函数，实际上是从令牌嵌入构建实体表示。它应该是一个将一组令牌（锚点、关系或两者）映射到单个向量的函数：

\[f([a_1, a_2, ...., a_k, r_1, r_2, ..., r_m]) \in \mathbb{R}^{(k+m) \times d} \rightarrow \mathbb{R}^{d}\]

目前，默认情况下我们使用一个简单的2层MLP (pykeen.nn.perceptron.ConcatMLP)，它将所有标记连接成一个长向量并将其投影到模型的嵌入维度：

hidden_dim = int(ratio * embedding_dim)
super().__init__(
    nn.Linear(num_tokens * embedding_dim, hidden_dim),
    nn.Dropout(dropout),
    nn.ReLU(),
    nn.Linear(hidden_dim, embedding_dim),
)

聚合可以使用任何神经网络进行参数化 (torch.nn.Module)，该网络可以从一组输入中返回单个向量。让我们变得花哨一点 😎 并创建一个 DeepSet 编码器：

class DeepSet(torch.nn.Module):
    def __init__(self, hidden_dim=64):
        super().__init__()
        self.encoder = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim, hidden_dim),
        )
        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim, hidden_dim),
        )

    def forward(self, x, dim=-2):
        x = self.encoder(x).mean(dim)
        x = self.decoder(x)
        return x


model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["AnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],
    embedding_dim=64,
    interaction="rotate",
    relation_initializer="init_phases",
    relation_constrainer="complex_normalize",
    entity_initializer="xavier_uniform_",
    aggregation=DeepSet(hidden_dim=64),
)

我们甚至可以在这里放置一个带有池化的Transformer。唯一需要注意的是编码器的复杂性——我们发现pykeen.nn.perceptron.ConcatMLP在速度和最终性能之间取得了良好的平衡，尽管代价是对输入标记集不具有排列不变性。

聚合函数类似于GNNs中的函数。在当前的分词设置中，非参数的avg/min/max效果不佳，因此一些非线性特性肯定是有用的——因此选择MLP / DeepSets / Transformer作为聚合函数。

让我们将我们酷炫的NodePiece模型与40/40/20的度/PageRank/随机标记化、BFS搜索器和DeepSet聚合封装到一个管道中：

result = pipeline(
    dataset="fb15k237",
    dataset_kwargs=dict(
        create_inverse_triples=True,
    ),
    model=NodePiece,
    model_kwargs=dict(
        tokenizers=["AnchorTokenizer", "RelationTokenizer"],
        num_tokens=[20, 12],
        tokenizers_kwargs=[
            dict(
                selection="MixtureAnchorSelection",
                selection_kwargs=dict(
                    selections=["degree", "pagerank", "random"],
                    ratios=[0.4, 0.4, 0.2],
                    num_anchors=500,
                ),
                searcher="ScipySparse",
            ),
            dict(),  # empty dict for the RelationTokenizer - it doesn't need any kwargs
        ],
        embedding_dim=64,
        interaction="rotate",
        relation_initializer="init_phases",
        relation_constrainer="complex_normalize",
        entity_initializer="xavier_uniform_",
        aggregation=DeepSet(hidden_dim=64),
    ),
)

NodePiece + GNN

也可以在获得的NodePiece表示之上添加消息传递GNN，以进一步丰富节点状态 - 我们发现它在归纳LP任务中显示出更好的结果。我们已经使用pykeen.models.InductiveNodePieceGNN实现了这一点，它使用了一个2层的CompGCN编码器 - 请查看归纳链接预测教程。

使用METIS对大型图进行标记化

在超过100万个节点的整个图上挖掘锚点并运行标记化可能在计算上非常昂贵。由于NodePiece的固有局部性，即通过最近的锚点和关联关系进行标记化，我们建议使用图分区来减少标记化的时间和内存成本。通过图分区，可以在每个分区内独立执行锚点搜索和标记化，最后将所有结果合并到一个词汇表中。

我们设计了使用METIS的分区标记化策略，这是一种最小割图分区算法，在torch-sparse中有高效的实现。除了METIS，我们还利用torch-sparse提供了一种新的、更快的BFS过程，可以在GPU上运行。

主要的标记器类是 pykeen.nn.node_piece.MetisAnchorTokenizer。 你可以用它来代替普通的 AnchorTokenizer。 使用基于Metis的标记器，我们首先将输入的训练图划分为 k 个独立的分区，然后依次且独立地为每个分区运行锚点选择和锚点搜索。

您可以使用上述任何现有的锚点选择和锚点搜索策略，尽管对于较大的图，我们建议使用新的pykeen.nn.node_piece.SparseBFSSearcher 作为锚点搜索器——它实现了更快的稀疏矩阵乘法内核，并且可以在GPU上运行。与普通分词器的唯一区别在于，现在num_anchors 参数定义了每个分区将挖掘多少个锚点。

新的分词器有两个特殊参数：

  • num_partitions - 图形将被划分成的分区数量。 你可以预期METIS会产生大小大致相同的分区，例如， num_partitions=10 对于一个有100万个节点的图形，将产生10个分区， 每个分区大约有10万个节点。挖掘的锚点总数将是 num_partitions * num_anchors

  • device - 运行METIS的设备。它可以与运行AnchorSearcher的设备不同。我们发现device="cpu"在较大的图上运行更快，并且不需要有限的GPU内存，尽管你可以让设备自动解析或设置为device="cuda"以尝试在GPU上运行。

仍然建议使用pykeen.nn.node_piece.SparseBFSSearcher在GPU上运行大型图标记化，因为稀疏CUDA内核更高效。如果有GPU可用，默认情况下会自动使用它。

让我们使用新的分词器来处理包含500万节点和2000万边的Wikidata5M图。

from pykeen.datasets import Wikidata5M

dataset = Wikidata5M(create_inverse_triples=True)

model = NodePiece(
    triples_factory=dataset.training,
    tokenizers=["MetisAnchorTokenizer", "RelationTokenizer"],
    num_tokens=[20, 12],  # 20 anchors per node in for the Metis strategy
    embedding_dim=64,
    interaction="rotate",
    tokenizers_kwargs=[
        dict(
            num_partitions=20,  # each partition will be of about 5M / 20 = 250K nodes
            device="cpu",  # METIS on cpu tends to be faster
            selection="MixtureAnchorSelection",  # we can use any anchor selection strategy here
            selection_kwargs=dict(
                selections=['degree', 'random'],
                ratios=[0.5, 0.5],
                num_anchors=1000,  # overall, we will have 20 * 1000 = 20000 anchors
            ),
            searcher="SparseBFSSearcher",  # a new efficient anchor searcher
            searcher_kwargs=dict(
                max_iter=5  # each node will be tokenized with anchors in the 5-hop neighborhood
            )
        ),
        dict()
    ],
    aggregation="mlp"
)

# we can save the vocabulary of tokenized nodes
from pathlib import Path
model.entity_representations[0].base[0].save_assignment(Path("./anchors_assignment.pt"))

在一台拥有32 GB内存和32 GB GPU的机器上，处理Wikidata5M大约需要10分钟：

  • 在CPU上分区为20个集群大约需要3分钟；

  • ~ 7 分钟用于每个分区中的锚点选择和搜索

我的图需要多少个分区？

这在很大程度上取决于现有的硬件和内存，但作为一般规则，我们建议每个分区的大小小于500K节点