2020-04-08 20:08:14
出品 | 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