训练超大规模图模型,PyTorchBigGraph如何作到?

Facebook 提出了一种可高效训练包含数十亿节点和数万亿边的图模型的框架 BigGraph 并开源了其 PyTorch 实现。本文将解读它的创新之处,解析它能从大规模图网络高效提取知识的缘由。git

 

训练超大规模图模型,PyTorchBigGraph如何作到?

图(graph)是机器学习应用中最基本的数据结构之一。具体来讲,图嵌入方法是一种无监督学习方法,可以使用本地图结构来学习节点的表征。社交媒体预测、物联网模式检测或药物序列建模等主流场景中的训练数据能够很天然地表征为图结构。其中每一种场景均可以轻松获得具备数十亿相连节点的图。图结构很是丰富且具备与生俱来的导向能力,所以很是适合机器学习模型。尽管如此,图结构却很是复杂,难以进行大规模扩展应用。也所以,现代深度学习框架对大规模图数据结构的支持仍很是有限。github

Facebook 推出过一个框架 PyTorch BigGraph:https://github.com/facebookresearch/PyTorch-BigGraph,它能更快更轻松地为 PyTorch 模型中的超大图结构生成图嵌入。服务器

某种程度上讲,图结构可视为有标注训练数据集的一种替代,由于节点之间的链接可用于推理特定的关系。这种方法遵守无监督图嵌入方法的模式,它能够学习图中每一个节点的向量表征,其具体作法是优化节点对的嵌入,使得之间有边相连的节点对的嵌入比无边相连的节点对的嵌入更近。这相似于在文本上训练的 word2vec 的词嵌入的工做方式。网络

训练超大规模图模型,PyTorchBigGraph如何作到?

当应用于大型图结构时,大多数图嵌入方法的结果都至关局限。举个例子,若是一个模型有 20 亿个节点,每一个节点有 100 个嵌入参数(用浮点数表示),则光是存储这些参数就须要 800 GB 内存,所以不少标准方法都超过了典型商用服务器的内存容量。这是深度学习模型面临的一大挑战,也是 Facebook 开发 BigGraph 框架的缘由。数据结构

 

 

PyTorch BigGraph多线程

PyTorch BigGraph(PBG)的目标是扩展图嵌入模型,使其有能力处理包含数十亿节点和数万亿边的图。PBG 为何有能力作到这一点?由于它使用了四大基本构建模块:框架

  1. 图分区,这让模型没必要彻底载入到内存中。
  2. 在每台机器上的多线程计算
  3. 在多台机器上的分布式执行(可选),全部操做都在图上不相连的部分进行
  4. 分批负采样,当每条边 100 个负例时,可实现每台机器每秒处理超过 100 万条边。

经过将图结构分区为随机划分的 P 个分区,使得可将两个分区放入内存中,PBG 解决了传统图嵌入方法的一些短板。举个例子,若是一条边的起点在分区 p1,终点在分区 p2,则它会被放入 bucket (p1, p2)。而后,在同一模型中,根据源节点和目标节点将这些图节点划分到 P2 bucket。完成节点和边的分区以后,能够每次在一个 bucket 内执行训练。bucket (p1, p2) 的训练仅须要将分区 p1 和 p2 的嵌入存储到内存中。PBG 结构能保证 bucket 至少有一个以前已训练的嵌入分区。机器学习

训练超大规模图模型,PyTorchBigGraph如何作到?

PBG 的另外一大创新是训练机制的并行化和分布式。PBG 使用 PyTorch 自带的并行化机制实现了一种分布式训练模型,这用到了前面描述的模块分区结构。在这个模型中,各个机器会协调在不相交的 bucket 上进行训练。这会用到一个锁服务器(lock server),其负责将 bucket 分派给工做器(worker),从而尽量地减小不一样机器之间的通讯。每台机器均可以使用不一样的 bucket 并行地训练模型。分布式

训练超大规模图模型,PyTorchBigGraph如何作到?

在上图中,机器 2 中的 Trainer 模块向机器 1 上的锁服务器请求了一个 bucket,这会锁定该 bucket 的分区。而后该 trainer 会保存它再也不使用的全部分区并从共享分区服务器载入它须要的新分区,此时它能够将本身的旧分区释放回锁服务器。而后边会从一个共享文件系统载入,并在没有线程内同步的状况下在多个线程上进行训练。在一个单独的线程中,仅有少许共享参数会与一个共享参数服务器持续同步。模型检查点偶尔会从 trainer 写入到共享文件系统中。这个模型容许使用至多 P/2 台机器时,让一组 P 个 bucket 并行化。ide

 

PBG 一项不那么直接的创新是使用了分批负采样技术。传统的图嵌入模型会沿真正例边将随机的「错误」边构建成负训练样本。这能显著提高训练速度,由于仅有一小部分权重必须使用每一个新样本进行更新。可是,负例样本最终会为图的处理引入性能开销,并最终会经过随机的源或目标节点「损害」真正的边。PBG 引入了一种方法,即复用单批 N 个随机节点以获得 N 个训练边的受损负例样本。相比于其它嵌入方法,这项技术让咱们能以很低的计算成本在每条边对应的许多负例上进行训练。

要增长在大型图上的内存效率和计算资源,PBG 利用了单批 Bn 个采样的源或目标节点来构建多个负例。在典型的设置中,PBG 会从训练集取一批 B=1000 个正例,而后将其分为 50 条边一个的块。来自每一个块的目标(与源等效)嵌入会与从尾部实体类型均匀采样的 50 个嵌入相连。50 个正例与 200 个采样节点的外积等于 9900 个负例。

训练超大规模图模型,PyTorchBigGraph如何作到?

分批负采样方法可直接影响模型的训练速度。若是没有分批,训练的速度就与负例的数量成反比。分批训练可改善方程,获得稳定的训练速度。

训练超大规模图模型,PyTorchBigGraph如何作到?

Facebook 使用 LiveJournal、Twitter 数据和 YouTube 用户互动数据等不一样的数据集评估了 PBG。此外,PBG 还使用 Freebase 知识图谱进行了基准测试,该知识图谱包含超过 1.2 亿个节点和 27 亿条边。另外还使用 Freebase 的一个小子集 FB15k 进行了测试,FB15k 包含 15000 个节点和 600000 条边,常被用做多关系嵌入方法的基准。FB15k 实验代表 PBG 的表现与当前最佳的图嵌入模型相近。可是,当在完整的 Freebase 数据集上评估时,PBG 的内存消耗获得了 88% 的改善。

训练超大规模图模型,PyTorchBigGraph如何作到?

PBG 是首个可扩展的、能训练和处理包含数十亿节点和数万亿边的图数据的方法。PBG 的首个实现已经开源,将来应该还会有更有意思的贡献。

 

 

【编辑推荐】

  1. 智能物联网:将人工智能的力量带入物联网
  2. 解决数独问题用人工智能仍是量子计算?
  3. 人工智能的企业家来讲,这四个新的方向可能值得重视
  4. 机器学习中的标签泄漏介绍及其如何影响模型性能
  5. 将来,你是「透明的」!计算机经过人脑信号预测想法,还原人脑图像准确率高达83%