PyTorch-BigGraph (PBG) 是一个分布式系统，用于学习大型图的图嵌入，尤其是具有多达数十亿个实体和数万亿条边的大型 Web 交互图。

PBG 是在PyTorch-BigGraph: A Large-scale Graph Embedding Framework论文中介绍的，该论文在 2019 年的SysML 会议上发表。

PBG 通过摄取其边列表在输入图上进行训练，每个边由其源实体和目标实体以及可能的关系类型标识。它为每个实体输出一个特征向量（嵌入），试图将相邻实体放在向量空间中彼此靠近，同时将未连接的实体推开。因此，具有相似邻居分布的实体最终会在附近。

可以将每种关系类型配置为以不同的方式计算此“接近度分数”，并使用在训练期间学习的参数（如果有）。这允许在多个关系类型之间共享相同的底层实体嵌入。

其模型的通用性和可扩展性允许 PBG 从知识图嵌入文献中训练出许多模型，包括TransE、RESCAL、DistMult和ComplEx。

PBG 在设计时考虑了规模，并通过以下方式实现：

• 图分区，因此模型不必完全加载到内存中
• 每台机器上的多线程计算
• 跨多台机器的分布式执行（可选），所有机器同时在图的不相交部分上运行
• 批量负采样，允许处理 >100 万条边/秒/机器，每个边有 100 个负片

PBG不适用于在小图上使用奇异模型进行模型探索，例如图卷积网络、深度网络等。

要求

PBG 是用 Python（3.6 版或更高版本）编写的，并且依赖于PyTorch（至少 1.0 版）和一些其他库。

所有计算都在 CPU 上执行，因此建议使用大量内核。不需要 GPU。

在多台机器上运行时，它们需要能够以高带宽（推荐 10 Gbps 或更高）相互通信，并能够访问共享文件系统（用于检查点）。PBG 使用torch.distributed，它使用在 TCP 或 MPI 之上运行的 Gloo 包。

安装

要安装最新版本的 PBG，请运行：

pip install torchbiggraph

作为一种替代方案，可以改为从存储库安装开发版本。这可能具有更新的功能，但可能更不稳定。为此，请克隆存储库（或将其下载为存档），然后在顶级目录中运行：

pip install .

入门

通过运行以下命令（执行此脚本）可以轻松复制论文的结果：

torchbiggraph_example_fb15k

这将下载 Freebase 15k 知识库数据集，将其放入正确的格式，使用 ComplEx 模型对其进行训练，最后对学习的嵌入进行评估，计算 MRR 和其他应该与论文匹配的指标。另一个命令torchbiggraph_example_livejournal对 LiveJournal 交互图数据集执行相同的操作。

要了解如何使用 PBG，让我们了解一下 FB15k 脚本的作用。

下载数据

首先，它检索数据集并对其进行解包，获得一个包含三个边缘集的目录作为 TSV 文件，用于训练、验证和测试。

wget https://dl.fbaipublicfiles.com/starspace/fb15k.tgz -P data
tar xf data/fb15k.tgz -C data

这些文件的每一行都包含有关一条边的信息。使用制表符作为分隔符，将行分为包含源实体、关系类型和目标实体的标识符的列。例如：

/m/027rn	/location/country/form_of_government	/m/06cx9
/m/017dcd	/tv/tv_program/regular_cast./tv/regular_tv_appearance/actor	/m/06v8s0
/m/07s9rl0	/media_common/netflix_genre/titles	/m/0170z3
/m/01sl1q	/award/award_winner/awards_won./award/award_honor/award_winner	/m/044mz_
/m/0cnk2q	/soccer/football_team/current_roster./sports/sports_team_roster/position	/m/02nzb8

准备数据

然后，脚本将边缘列表转换为 PBG 的输入格式。这相当于为所有实体和关系类型分配一个数字标识符，对实体和边进行洗牌和分区，并以正确的格式写下来。

幸运的是，有一个命令可以完成所有这些操作：

torchbiggraph_import_from_tsv \
  --lhs-col=0 --rel-col=1 --rhs-col=2 \
  torchbiggraph/examples/configs/fb15k_config.py \
  data/FB15k/freebase_mtr100_mte100-train.txt \
  data/FB15k/freebase_mtr100_mte100-valid.txt \
  data/FB15k/freebase_mtr100_mte100-test.txt

输出将存储在data/FB15k目录中的输入旁边。

这个简单的实用程序只适用于完全适合内存的小图。要处理更大的数据，必须实现自己的自定义预处理器。

训练

该torchbiggraph_train命令用于启动训练。训练参数隐藏在配置文件中，该文件的路径提供给命令。但是，它们可以从带有--param标志的命令行覆盖。示例配置用于训练和评估，因此我们必须使用覆盖来指定要使用的边缘集。

torchbiggraph_train \
  torchbiggraph/examples/configs/fb15k_config.py \
  -p edge_paths=data/FB15k/freebase_mtr100_mte100-train_partitioned

这将从entity_path配置中指定的edge_paths目录和命令行中给出的目录中读取数据。它将检查点（也作为输出数据的两倍）写入checkpoint_path也在配置中定义的目录，在这种情况下是model/fb15k.

训练总共将进行 50 个 epoch，并将进度和一些统计信息记录到控制台，例如：

Starting epoch 1 / 50, edge path 1 / 1, edge chunk 1 / 1
Edge path: data/FB15k/freebase_mtr100_mte100-train_partitioned
still in queue: 0
Swapping partitioned embeddings None ( 0 , 0 )
( 0 , 0 ): Loading entities
( 0 , 0 ): bucket 1 / 1 : Processed 483142 edges in 17.36 s ( 0.028 M/sec ); io: 0.02 s ( 542.52 MB/sec )
( 0 , 0 ): loss:  309.695 , violators_lhs:  171.846 , violators_rhs:  165.525 , count:  483142
Swapping partitioned embeddings ( 0 , 0 ) None
Writing partitioned embeddings
Finished epoch 1 / 50, edge path 1 / 1, edge chunk 1 / 1
Writing the metadata
Writing the checkpoint
Switching to the new checkpoint version

评估

训练完成后，它产生的实体嵌入可以根据保留的边缘集进行评估。该torchbiggraph_example_fb15k命令执行过滤评估，通过将评估集中的边缘与所有其他边缘进行比较来计算评估集中的等级，除了在任何训练、验证或测试集中为真阳性的边缘。FB15k 的文献中使用了过滤评估，但不能扩展到小图之外。

最终结果应与论文mrr中报告的(Mean Reciprocal Rank, MRR) 和r10(Hits@10)的值相匹配：

Stats: pos_rank:  65.4821 , mrr:  0.789921 , r1:  0.738501 , r10:  0.876894 , r50:  0.92647 , auc:  0.989868 , count:  59071

评估也可以直接从命令行运行，如下所示：

torchbiggraph_eval \
  torchbiggraph/examples/configs/fb15k_config.py \
  -p edge_paths=data/FB15k/freebase_mtr100_mte100-test_partitioned \
  -p relations.0.all_negs=true \
  -p num_uniform_negs=0

但是，无法在命令行上进行过滤评估，因此报告的结果将与论文不符。它们将类似于：

Stats: pos_rank:  234.136 , mrr:  0.239957 , r1:  0.131757 , r10:  0.485382 , r50:  0.712693 , auc:  0.989648 , count:  59071

转换输出

在预处理期间，实体和关系类型的标识符从字符串转换为序数。为了将输出嵌入映射回原始名称，可以执行以下操作：

torchbiggraph_export_to_tsv \
  torchbiggraph/examples/configs/fb15k_config.py \
  --entities-output entity_embeddings.tsv \
  --relation-types-output relation_types_parameters.tsv

这将创建entity_embeddings.tsv文件，这是一个文本文件，其中每一行包含实体的标识符，后跟其嵌入的组件，每个组件位于不同的列中，全部由制表符分隔。例如，为简洁起见缩短每一行：

/m/0fphf3v	-0.524391472	-0.016430536	-0.461346656	-0.394277513	0.125605106	...
/m/01bns_	-0.122734159	-0.091636233	0.506501377	-0.503864646	0.215775326	...
/m/02ryvsw	-0.107151665	0.002058491	-0.094485454	-0.129078045	-0.123694092	...
/m/04y6_qr	-0.577532947	-0.215747222	-0.022358289	-0.352154016	-0.051905245	...
/m/02wrhj	-0.593656778	-0.557167351	0.042525314	-0.104738958	-0.265990764	...

它还将创建一个relation_types_parameters.tsv文件，其中包含关系类型的运算符参数。格式与上面类似，但每一行以更多的键列开始，分别包含关系类型的名称、边（lhs或rhs）、该边上该关系类型使用的运算符的名称、该运算符的参数名称和参数的形状（整数由 分隔x）。这些列后面是 flattened 参数的值。例如，对于两种关系类型foo和bar，分别使用运算符linear和complex_diagonal，嵌入维度为 200 并启用动态关系，此文件可能如下所示：

foo	lhs	linear	linear_transformation	200x200	-0.683401227	0.209822774	-0.047136042	...
foo	rhs	linear	linear_transformation	200x200	-0.695254087	0.502532542	-0.131654695	...
bar	lhs	complex_diagonal	real	200	0.263731539	1.350529909	1.217602968	...
bar	lhs	complex_diagonal	imag	200	-0.089371338	-0.092713356	0.025076168	...
bar	rhs	complex_diagonal	real	200	-2.350617170	0.529571176	0.521403074	...
bar	rhs	complex_diagonal	imag	200	0.692483306	0.446569800	0.235914066	...

文档

更多信息可以在完整文档中找到。

预训练嵌入

我们在完整的Wikidata图上训练了一个 PBG 模型，使用翻译算子来表示关系。它可以在这里下载（36GiB，gzip 压缩）。我们使用此处数据的真实版本来训练我们的模型。如上节所述，模型文件采用 TSV 格式。请注意，文件的第一行包含实体的数量、关系的数量和嵌入的维度，以制表符分隔。该模型包含 7800 万个实体、4131 个关系，嵌入的维度为 200。

引文

要引用这项工作，请使用：

@inproceedings{pbg,
  title={{PyTorch-BigGraph: A Large-scale Graph Embedding System}},
  author={Lerer, Adam and Wu, Ledell and Shen, Jiajun and Lacroix, Timothee and Wehrstedt, Luca and Bose, Abhijit and Peysakhovich, Alex},
  booktitle={Proceedings of the 2nd SysML Conference},
  year={2019},
  address={Palo Alto, CA, USA}
}

执照

PyTorch-BigGraph 已获得 BSD 许可，可在LICENSE.txt文件中找到。