节点嵌入训练加快300倍:开源高性能图嵌入系统GraphVite技术解读

GraphVite是一个专为高速、大规模的训练场景设计的通用图嵌入系统,已于本周正式开源。该系统旨在提供一个包含一系列图嵌入方法的通用高性能框架,进一步推动图学习算法的研究与发展。GraphVite支持多GPU并行,可以扩展到百万级甚至十亿级的图,在百万节点的图上,只需要约一分钟就能完成节点表示的学习。GraphVite由唐建教授带领的MilaGraph团队开发,相关论文已发表于WWW19会议。本文是AI前线第88篇论文导读,我们将对GraphVite及其技术实现进行具体介绍。

image

GraphVite 总览

目前,GraphVite已为3种应用提供了完整的训练和评价框架,包括节点嵌入、知识图谱嵌入以及图和高维数据可视化。此外,它还包含了9种流行的模型以及这些模型在一系列标准数据集上的基准测试结果。

GitHub项目入口:https://github.com/DeepGraphLearning/graphvite

image

GraphVite现支持的三种任务

为了展示GaphVite的速度,MilaGraph提供了在三种应用上最好的开源实现与GraphVite的训练速度报告,报告基于24个CPU线程以及4个Tesla V100GPU的硬件系统。在Youtube数据集上,节点嵌入任务的训练速度如下:

image

Youtube数据集上节点嵌入任务训练时间

在FB15k数据库上实现的知识图谱任务报告如下:

image

在MNIST数据集上实现的高维数据可视化任务报告如下:

image

可以看出在以上三个任务中,相较于目前最好的开源实现,GraphVite最多可以提升300多倍的速度。并且,在提高了速度的同时,GraphVite没有任何性能上的牺牲,这点可以从该项目提供的基准测试报告中发现。在使用与上述速度报告相同硬件的情况下,GraphVite的说明文档中还提供了使用GraphVite实现的所有模型的基准测试结果。在节点嵌入任务上,作者在Youtube、Flicke和Friendster-small三个数据集上提供了每个模型的micro-F1和macro-F1指标。

image

Youtube数据集上三种不同模型的结果报告

在知识图谱嵌入任务上,作者对TransE、DistMult、ComplE、SimplE和RotatE四种方法进行了测试。

image

Youtube数据集上五种不同模型的结果报告

作者在两个流行的数据集上对图及高维数据可视化任务进行了测试。LargeVis所需的训练时间和资源如下表所示。请注意,超过95%的GPU内存成本来自KNN Graph的构建,并且可以在必要时与速度进行权衡。

image

在MNIST和ImageNet数据集上LargeVis的测试结果

其中,MNIST数据集上3D可视化结果如下:

image

与类似的工作Pytorch-BigGraph相比,GraphVite速度更快,测试结果如下:

image

同时,GraphVite还比前者多了节点嵌入以及可视化两个功能,为研究和开发提供了更多的便利。

理解GraphVite

GraphVite为什么可以这么快?这是因为GraphVite根据CPU和GPU各自的优势,解决了混合节点嵌入系统的三大挑战。

混合节点嵌入系统的难点主要包含:1)基于mini-batch的SGD方法无法直接用关于大型网络上的节点嵌入,2)由于GPU与CPU总线带宽的差异,在CPU与GPU之间的数据转送会成为系统瓶颈,3) 多CPU与GPU设备间的同步成本巨大。

GraphVite使用了并行在线增广方法在CPU上有效地扩充网络规模。同时使用了并行负采样方法对多GPU的嵌入训练进行协调。GraphVite还使用了一种协作策略来降低CPU和GPU之间的同步成本。

image

GraphVite系统结构图,灰色及黄色部分分别表示并行数据增广过程与嵌入训练过程。

并行在线增广

对于节点嵌入方法,其第一阶段需要使用随机遍历来扩充原始网络。由于增广网络通常比原始网络要大一到两个数量级,因此,如果原始网络已经非常大,则其网络结构将很难被加载到主存储器中。GraphVite采用了并行在线增广,可以在不显示网络增强的情况下生成增广的边样本。

作者首先绘制一个出发节点(depature node),其选择概率与每个节点的度成正比。然后从出发节点开始执行随机游走,并在特定的增广距离s内挑选节点对产生边样本。由于随机游走中生成边样本具有相关性,因此训练效果可能会变差。作者受到强化学习中广泛使用的经验回放方法的启发,收集边样本到样本池,并在对它混洗(shuffle)后再转发到GPU进行嵌入训练。当每个线程预先分配有独立的样池时,上述方法可以并行化处理。算法2给出了并行在线增广的详细过程:

image

为了更好的训练效果,对样本池进行混洗非常重要,但同时它也会减慢网络增广的速度。原因是普通的洗牌方法包含大量随机内存访问,并且不能由CPU缓存加速。如果服务器有多个CPU插槽,速度损失将更加严重。为了缓解这个问题,作者提出了一种伪混洗技术,它以更加缓存友好的方式对相关样本进行洗牌,并显着提高系统的速度。

并行负采样

在嵌入训练阶段,作者将训练任务分为多个子任务并将它们分配给多个GPU设备。子任务必须设计成具有很少的共享数据的形式以最小化GPU之间的同步成本。为了了解如何将模型参数分配给多个GPU而不重叠,作者首先介绍了ε-梯度可交换的概念:

image

其中,作者将0-梯度可交换称为梯度可交换。由于节点嵌入训练过程的稀疏性质,在网络中会存在很多形成梯度可交换对的集合。例如,对于两个边缘样本集例如,对于两个边缘样本集X1,X2⊆E,如果它们不共享任何源节点或目标节点,则X1和X2是梯度可交换的。即使X1和X2共享一些节点,如果学习率α和迭代次数有界,它们仍然可以是ε-梯度可交换的。

image

基于节点嵌入中观察到的梯度交换性,作者提出了一种用于嵌入训练阶段的并行负采样算法。对于n个GPU,我们将顶点和上下文行分别划分为n个分区(上图的左上角)。这导致样本池的n×n分区网格,其中每个边缘属于其中一个块。这样,任何不共享行或列的块都是可梯度交换的。只要限制每个块上的迭代次数,同一行或列中的块都是ε-梯度可交换的。

作者将集(episode)定义为并行负抽样中使用的块级步骤。在每一集中,我们分别向n个GPU发送n个正交块及其对应的顶点和上下文分区。然后,每个GPU使用ASGD更新自己的嵌入分区。因为这些块是相互梯度可交换的并且不共享参数矩阵中的任何行,所以多个GPU可以在没有任何同步的情况下同时执行ASGD。在每集结束时,我们从所有GPU收集更新的参数并分配另外n个正交块。这里ε-梯度可交换由n个正交块中的总样本数控制,我们将其定义为集的大小。较小的集有利于ε-梯度可交换嵌入式训练。但是同时也会导致更频繁的同步。因此作者调整了集的大小,以便在速度和ε-梯度可交换之间有一个很好的权衡。上图给出了具有4个分区的并行负采样的示例。

虽然作者使用等于n的分区数量来说明了负采样的有效性,但是并行负采样可以很容易地推广到任何数量大于n的分区的情况,只需在每集中处理n个子组中的正交块。算法3给出了用于多个GPU的混合系统的并行负采样算法的详细步骤:

image

协作策略

我们的并行负采样使不同的GPU能够同时训练节点嵌入,只需要在集之间进行同步。但是,应该注意到,样本池也在CPU和GPU之间共享。如果它们在样本池上同步,则只有同一阶段的设备可以同时访问样本池,这意味着硬件在一半时间内处于空闲状态。为了提高效率,作者提出了一种协作策略来降低同步成本。

作者主存储器中分配两个样本池,让CPU和GPU始终在不同的样本池上工作。 CPU首先填充样本池并将其传递给GPU。之后,并行在线增广和并行负采样分别在CPU和GPU上执行。当CPU填满新的样本池时,将交换这两个池。通过协作策略,CPU和GPU之间的同步成本降低,混合系统的速度几乎翻倍。

实验

作者使用了Youtube、Reiendster-small、Hyperlink-PLD和Friendster四个数据库对GraphVite的性能进行了测试。部分实验结果已在第一章中展示,详细的实验结果在论文原文以及GraphVite官方网站https://graphvite.io/)的Docs中都有详细的展示。

快速入门

GraphVite可以在CUDA>=9.2的Linux系统上运行,兼容Python2.7和Python3.5/3.6/3.7版本。GraphVite对待新手十分友好,它提供了关于节点嵌入任务的快速入门示例供用户学习。用户可以选择通过Conda源进行安装或使用源码编译,使用Conda安装仅需在命令行输入命令:

conda install -c milagraph graphvite

如果仅需要嵌入训练而不要测试部分,还可以选择更加轻量级的最小化GraphVite系统:

conda install -c milagraph graphvite-mini

或者用户也可以选择使用源码编译:

git clone https://github.com/DeepGraphLearning/graphvite
cd graphvite
conda install -y --file conda/requirements.txt
mkdir build
cd build && cmake .. && make && cd -
cd python && python setup.py install && cd -

安装完成GraphVite后,直接执行即可查看示例,在终端输入:

conda install -c milagraph graphvite-mini

示例需要运行约一分多钟的时间,然后会输出结果:

Batch id: 6000
loss = 0.371641

macro-F1@20%: 0.236794
micro-F1@20%: 0.388110
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章