半小时演习亿级规模常识图谱亚马逊这个 AI 框架要火

知识图谱 （Knowledge Graph）作为一个主要的技能，在近几年里被广泛利用在了信息检索，自然措辞处理，以及推举系统等各种领域。
学习知识图谱的嵌入表示 （Knowledge Graph Embeddings）是一种从知识图谱构造中天生无监督节点特色（node feature）的方法，天生的特色可以被用在各种机器学习任务之上。
例如，可以通过节点的嵌入表示来预测两个节点之间是否有链接（link prediction）。

然而，随着社交网络、推举系统等范例图数据场景的发展，知识图谱的规模也在不断地增长。
在工业界真实的场景中，技能职员常常须要面对千万级，乃至是亿万级节点的大规模图数据。
如何快速、高效地在大规模知识图谱上进行嵌入表示的演习是当前的一个寻衅。

近日，亚马逊 AI 团队继 DGL 之后，又开源了一款专门针对大规模知识图谱嵌入表示的新演习框架 DGL-KE，旨在能让研究职员和工业界用户方便、快速地在大规模知识图谱数据集上进行机器学习演习任务。

github 地址：https://github.com/awslabs/dgl-ke

比较于已有的开源框架，DGL-KE 的亮点如下：

支持各种主流知识图谱表示学习算法，包括 TransE、ComplEx、DistMult、TransR、RESCAL、RotatE 等；

已有开源框架中唯一支持多核 CPU，多卡 GPU，CPU-GPU 稠浊演习，以及分布式演习的知识图谱嵌入表示框架；

大略易用，用户无需编写代码，直接将知识图谱数据作为输入即可；

高性能且可扩展。
根据 DGL-KE 发布的 Freebase 数据集 （超过 8600 万节点，3 亿条边）Benchmark 显示，在 AWS EC2 平台上，一台 p3.16xlarge （8 GPUs）可以在 100 分钟内完成演习。
4 台 r5dn.24xlarge （448 CPUs）可以在 30 分钟内完成演习，并且达到线性加速比。
这一结果比当前最快的同类系统（如 Facebook 发布 Pytorch-BigGraph）快 2-5 倍。

图1: DGL-KE 系统架构

DGL-KE 之以是能够有这样的性能，紧张是由于采取了许多创新的系统和算法优化：

（一）基于 METIS 图分割算法的分布式演习

对超大规模的图数据进行演习，分布式演习必不可少。
其思路紧张是将原始大图分割身分歧的子图，每一台机器卖力在一个子图上进行随机梯度低落演习，所有机器之间通过参数做事器（Parameter Server）进行模型的同步。
其架构如下图所示：

图2: DGL 分布式架构

然而，如果只是对一张大图进行随机切割，会造成演习机器与参数做事器之间的数据通信量巨大（本地机器须要从远程机器去要求自己所需的模型数据），从而造成网络瓶颈。
为理解决这一问题，DGL-KE 在演习前会预先通过 METIS 图分割算法对原始数据进行切割。

METIS 算法是打算机科学家 George Karypis 于 1995 年提出的一种高效的图分割算法，而 George Karypis 也正是 DGL-KE 项目的作者之一。
METIS 算法会将一张大图上干系联的节点尽可能放置在同一个划分（partition）之中。
这样可以将大部份的网络通信开销转化成本地机器的内存拷贝，从而大大提升了分布式演习的速率。

在 Freebase 这个数据集的实际演习中，METIS 算法可以节省将近 90% 的模型网络传输带宽，从而使分布式演习达到线性加速比。
DGL-KE 的分布式演习利用了 DGL-KVStore 组件。
DGL-KVStore 是专门为 DGL 系统定制开拓的参数做事器模块，用来实现稀疏模型的同步通讯。
该组件通过 C++ 实现了底层 socket、行列步队，以及稀疏数据序列化的定向优化，并且可以无缝兼容 METIS 图分割算法。

（二）基于共享内存的单机多进程演习

多核（Multi-core）已经成为当前打算机体系架构的标配，很多强大的 workstation 在一台机器内更是会有超过几十个 CPU 核心和上百 GB 乃至上 T 的内存。
对付很多千万级节点的图数据来说，这样的单机性能已经足够处理这种规模的数据。

DGL-KE 针对这样的场景也做了相应的系统优化，让用户尽可能地挖掘一台机器的性能极限。
与传统基于多线程（Multi-thread）的并行优化不同，DGL-KE 采取了基于多进程（Multi-Process）的粗粒度并行优化。
粗粒度的并行可以最大限的提升程序运行并行度，从而提高加速比。
此外，DGL-KE 在不同进程之间通过共享内存（Shared-memory）进行模型同步，从而大大减小了进程之间的通信开销。

图3: 基于共享内存的单机多进程演习

（三）CPU-GPU 稠浊演习

知识图谱嵌入表示的演习过程中会产生大量的矩阵运算，而矩阵运算可以通过 GPU 来加速。
对付小规模的图数据来说，DGL-KE 许可用户将完全的图模型放入 GPU 中进行演习，从而达到最优性能。
然而，比较于 CPU 内存，GPU 内存要小很多，一但模型 Embeddings 的大小超过了 GPU 内存限定就无法进行演习。
针对这样的场景，DGL-KE 为用户供应了 CPU-GPU 稠浊演习的模式。

在 CPU-GPU 稠浊演习模式中，模型 Embeddings 被存储在 CPU 的内存里，而 GPU 通过 mini-batch 的办法在每一轮迭代过程中将一小部分数据从 CPU 拷贝到 GPU 进行演习。
为了避免 CPU 与 GPU 之间的数据拷贝开销，DGL-KE 采取异步演习的办法将数据拷贝与打算 overlap 起来。
然而，异步打算会带来模型收敛速率和精确度的低落，DGL-KE 在这里采取了另一个优化，将 Entity Embedding 和 Relation Embedding 的更新采取不同的办法进行：relation 利用同步更新，而 enity 利用异步更新。

之以是这样做是由于在实际的演习过程中，relation 在很多数据集上都表现为长尾分布，即某几种 relation 类型霸占绝大多数，以是异步更新会导致 relation embedding 在演习过程中产生大量的模型冲突，从而影响模型的收敛和准确性。
而 entity 在演习过程常日是稀疏的，以是异步演习只会产生很小的冲突。
采取这样一个大略的优化，DGL-KE 既可以担保模型演习的收敛性，又可以担保系统性能。

图4: CPU-GPU 稠浊演习

除了以上优化之外，DGL-KE 还供应了其他多少优化方法。
例如，利用 Joint Negative Sampler 加速负采样过程，利用 Relation Partition 来减少演习过程中的数据拷贝，以及利用 Periodic synchronization 担保模型的收敛等。
DGL-KE 内置了多个处理好格式的知识图谱数据集，用户可以直接下载利用。

除此之外，DGL-KE 在两个小数据集 FB15k，wn18，以及一个大数据集 Freebase 上供应了演习 Benchmark，用户可以直接通过供应的脚本复现演习结果。
和已有的开源框架比较，DGL-KE 的性能上风明显，如下结果展示了 DGL-KE 与 Graphvite 在 FB15k 数据集上的性能比较，以及和 Pytorch-Biggraph 在 Freebase 数据集上的比较。

DGL-KE vs Graphvite

DGL-KE vs Pytorch-Biggraph

☞AI 四巨子 Google、DeepMind、Microsoft、Uber 深度学习框架大比拼

☞腾讯口试官这样问我二叉树，我刚好都会 | 原力操持

☞斩获GitHub 2000+ Star，阿里云开源的 Alink 机器学习平台如何跑赢双11数据“博弈”？| AI 技能生态论

☞微软为一人收购一公司？破解索尼程序、写黑客小说，看他彪悍的程序人生！

☞机器学习项目模板：ML项目的6个基本步骤

☞IBM、微软、苹果、谷歌、三星……这些区块链中的科技巨子原来已经做了这么多事！

☞资深程序员总结：剖析Linux进程的6个方法，我全都见告你

你点的每一个赞，我都负责当成了喜好