[极速聊科技]亚马逊AI开源知识图谱嵌入表示框架DGL-KE，半小时训练亿级规模知识图谱知识图谱作为一个重要的技术

知识图谱作为一个重要的技术，在近几年里被广泛运用在了信息检索，自然语言处理，以及推荐系统等各种领域。学习知识图谱的嵌入表示是一种从知识图谱结构中生成无监督节点特征的方法，生成的特征可以被用在各种机器学习任务之上。例如，可以通过节点的嵌入表示来预测两个节点之间是否有链接。
然而，随着社交网络、推荐系统等典型图数据场景的发展，知识图谱的规模也在不断地增长。在工业界真实的场景中，技术人员常常需要面对千万级，甚至是亿万级节点的大规模图数据。如何快速、高效地在大规模知识图谱上进行嵌入表示的训练是当前的一个挑战。
近日，亚马逊AI团队继DGL之后，又开源了一款专门针对大规模知识图谱嵌入表示的新训练框架DGL-KE ，旨在能让研究人员和工业界用户方便、快速地在大规模知识图谱数据集上进行机器学习训练任务。
相比于已有的开源框架， DGL-KE的亮点如下：
支持各种主流知识图谱表示学习算法，包括TransE、ComplEx、DistMult、TransR、RESCAL、RotatE等；
已有开源框架中唯一支持多核CPU ，多卡GPU ， CPU-GPU混合训练，以及分布式训练的知识图谱嵌入表示框架；
简单易用，用户无需编写代码，直接将知识图谱数据作为输入即可；
高性能且可扩展。根据DGL-KE发布的Freebase数据集Benchmark显示，在AWSEC2平台上，一台p3.16xlarge可以在100分钟内完成训练。 4台r5dn.24xlarge可以在30分钟内完成训练，并且达到线性加速比。这一结果比当前最快的同类系统快2-5倍。

文章图片
【[极速聊科技]亚马逊AI开源知识图谱嵌入表示框架DGL-KE，半小时训练亿级规模知识图谱】图1:DGL-KE系统架构
DGL-KE之所以能够有这样的性能，主要是因为采用了许多创新的系统和算法优化：
基于METIS图分割算法的分布式训练
对超大规模的图数据进行训练，分布式训练必不可少。其思路主要是将原始大图分割成不同的子图，每一台机器负责在一个子图上进行随机梯度下降训练，所有机器之间通过参数服务器进行模型的同步。其架构如下图所示：

文章图片
图2:DGL分布式架构
然而，如果只是对一张大图进行随机切割，会造成训练机器与参数服务器之间的数据通信量巨大，从而造成网络瓶颈。为了解决这一问题， DGL-KE在训练前会预先通过METIS图分割算法对原始数据进行切割。
METIS算法是计算机科学家GeorgeKarypis于1995年提出的一种高效的图分割算法，而GeorgeKarypis也正是DGL-KE项目的作者之一。 METIS算法会将一张大图上相关联的节点尽可能放置在同一个划分之中。这样可以将大部份的网络通信开销转化成本地机器的内存拷贝，从而大大提升了分布式训练的速度。
在Freebase这个数据集的实际训练中， METIS算法可以节省将近90%的模型网络传输带宽，从而使分布式训练达到线性加速比。 DGL-KE的分布式训练使用了DGL-KVStore组件。 DGL-KVStore是专门为DGL系统定制开发的参数服务器模块，用来实现稀疏模型的同步通讯。该组件通过C++实现了底层socket、消息队列，以及稀疏数据序列化的定向优化，并且可以无缝兼容METIS图分割算法。
基于共享内存的单机多进程训练
多核已经成为当前计算机体系架构的标配，很多强大的workstation在一台机器内更是会有超过几十个CPU核心和上百GB甚至上T的内存。对于很多千万级节点的图数据来说，这样的单机性能已经足够处理这种规模的数据。
DGL-KE针对这样的场景也做了相应的系统优化，让用户尽可能地挖掘一台机器的性能极限。与传统基于多线程的并行优化不同， DGL-KE采用了基于多进程的粗粒度并行优化。粗粒度的并行可以最大限的提升程序运行并行度，从而提高加速比。此外， DGL-KE在不同进程之间通过共享内存进行模型同步，从而大大减小了进程之间的通信开销。