半小时训练亿级规模知识图谱，亚马逊AI开源知识图谱嵌入表示框架DGL-KE
出品 | AI科技大本营（ID:rgznai100）
知识图谱 （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 （4*48 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大调研》，只需2分钟，便可收获价值299元的「AI开发者万人大会」在线直播门票!
你点的每个“在看”，我都认真当成了AI