PyTorch-BigGraph (PBG) 是一个用于学习大型图嵌入的分布式系统，尤其适用于包含数十亿实体和万亿边的大型网页交互图。

PBG 在 2019 年 SysML 会议的论文《PyTorch-BigGraph: A Large-scale Graph Embedding Framework》中被提出。

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 个训练周期，进度和一些统计信息将记录到控制台，例如

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（平均倒数排名，MRR）和 r10（命中@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 分隔的整数）。这些列后面跟着展开参数的值。例如，对于两个关系类型，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文件中所示。