半小时训练亿级规模知识图谱，亚马逊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在不同进程之间通过共享内存进行模型同步 ， 从而大大减小了进程之间的通信开销 。
文章图片
图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在这里采用了另一个优化 ， 将EntityEmbedding和RelationEmbedding的更新采用不同的方式进行：relation使用同步更新 ， 而enity使用异步更新 。
之所以这样做是因为在实际的训练过程中 ， relation在很多数据集上都表现为长尾分布 ， 即某几种relation类型占据绝大多数 ， 所以异步更新会导致relationembedding在训练过程中产生大量的模型冲突 ， 从而影响模型的收敛和准确性 。 而entity在训练过程通常是稀疏的 ， 所以异步训练只会产生很小的冲突 。 采用这样一个简单的优化 ， DGL-KE既可以保证模型训练的收敛性 ， 又可以保证系统性能 。
文章图片
图4:CPU-GPU混合训练
除了以上优化之外 ， DGL-KE还提供了其他若干优化方法 。 例如 ， 使用JointNegativeSampler加速负采样过程 ， 使用RelationPartition来减少训练过程中的数据拷贝 ， 以及使用Periodicsynchronization保证模型的收敛等 。 DGL-KE内置了多个处理好格式的知识图谱数据集 ， 用户可以直接下载使用 。
除此之外 ， DGL-KE在两个小数据集FB15k ， wn18 ， 以及一个大数据集Freebase上提供了训练Benchmark ， 用户可以直接通过提供的脚本复现训练结果 。 和已有的开源框架相比 ， DGL-KE的性能优势明显 ， 如下结果展示了DGL-KE与Graphvite在FB15k数据集上的性能比较 ， 以及和Pytorch-Biggraph在Freebase数据集上的比较 。
DGL-KEvsGraphvite
DGL-KEvsPytorch-Biggraph
推荐阅读
- 整挺好！门迪：当你回来训练时，发型好心情自然好
- 早期苏联炮兵实弹训练照片：装备老式火炮 指挥官腰间悬挂马刀
- 儿萨梦？阿贾克斯小将德斯特身穿巴萨训练服出镜
- 每体：阿贾克斯小将德斯特穿巴萨训练服 他已决定加盟巴萨
- 多图流：曼城恢复训练，萨内回归，瓜帅全副武装
- “我们每天都要竭尽全力地训练”
- 双外跟队训练，三名大将表决心，马布里带领北控冲击前四
- 迷你罗上线，C罗今日带儿子到尤文训练基地训练
- 半小时跑不出10公里，速度不如自行车，新国标电动车为何不提速？
- 了不起！于大宝公开力挺李铁，直言他让国脚在训练中都变得非常拼