GraphVite是一个专为高速、大规模的训练场景设计的通用图嵌入系统,已于本周正式开源。该系统旨在提供一个包含一系列图嵌入方法的通用高性能框架,进一步推动图学习算法的研究与发展。GraphVite支持多GPU并行,可以扩展到百万级甚至十亿级的图,在百万节点的图上,只需要约一分钟就能完成节点表示的学习。GraphVite由唐建教授带领的MilaGraph团队开发,相关论文已发表于WWW19会议。本文是AI前线第88篇论文导读,我们将对GraphVite及其技术实现进行具体介绍。
GraphVite 总览
目前,GraphVite已为3种应用提供了完整的训练和评价框架,包括节点嵌入、知识图谱嵌入以及图和高维数据可视化。此外,它还包含了9种流行的模型以及这些模型在一系列标准数据集上的基准测试结果。
GitHub项目入口:https://github.com/DeepGraphLearning/graphvite
为了展示GaphVite的速度,MilaGraph提供了在三种应用上最好的开源实现与GraphVite的训练速度报告,报告基于24个CPU线程以及4个Tesla V100GPU的硬件系统。在Youtube数据集上,节点嵌入任务的训练速度如下:
在FB15k数据库上实现的知识图谱任务报告如下:
在MNIST数据集上实现的高维数据可视化任务报告如下:
可以看出在以上三个任务中,相较于目前最好的开源实现,GraphVite最多可以提升300多倍的速度。并且,在提高了速度的同时,GraphVite没有任何性能上的牺牲,这点可以从该项目提供的基准测试报告中发现。在使用与上述速度报告相同硬件的情况下,GraphVite的说明文档中还提供了使用GraphVite实现的所有模型的基准测试结果。在节点嵌入任务上,作者在Youtube、Flicke和Friendster-small三个数据集上提供了每个模型的micro-F1和macro-F1指标。
在知识图谱嵌入任务上,作者对TransE、DistMult、ComplE、SimplE和RotatE四种方法进行了测试。
作者在两个流行的数据集上对图及高维数据可视化任务进行了测试。LargeVis所需的训练时间和资源如下表所示。请注意,超过95%的GPU内存成本来自KNN Graph的构建,并且可以在必要时与速度进行权衡。
其中,MNIST数据集上3D可视化结果如下:
与类似的工作Pytorch-BigGraph相比,GraphVite速度更快,测试结果如下:
同时,GraphVite还比前者多了节点嵌入以及可视化两个功能,为研究和开发提供了更多的便利。
理解GraphVite
GraphVite为什么可以这么快?这是因为GraphVite根据CPU和GPU各自的优势,解决了混合节点嵌入系统的三大挑战。
混合节点嵌入系统的难点主要包含:1)基于mini-batch的SGD方法无法直接用关于大型网络上的节点嵌入,2)由于GPU与CPU总线带宽的差异,在CPU与GPU之间的数据转送会成为系统瓶颈,3) 多CPU与GPU设备间的同步成本巨大。
GraphVite使用了并行在线增广方法在CPU上有效地扩充网络规模。同时使用了并行负采样方法对多GPU的嵌入训练进行协调。GraphVite还使用了一种协作策略来降低CPU和GPU之间的同步成本。
并行在线增广
对于节点嵌入方法,其第一阶段需要使用随机遍历来扩充原始网络。由于增广网络通常比原始网络要大一到两个数量级,因此,如果原始网络已经非常大,则其网络结构将很难被加载到主存储器中。GraphVite采用了并行在线增广,可以在不显示网络增强的情况下生成增广的边样本。
作者首先绘制一个出发节点(depature node),其选择概率与每个节点的度成正比。然后从出发节点开始执行随机游走,并在特定的增广距离s内挑选节点对产生边样本。由于随机游走中生成边样本具有相关性,因此训练效果可能会变差。作者受到强化学习中广泛使用的经验回放方法的启发,收集边样本到样本池,并在对它混洗(shuffle)后再转发到GPU进行嵌入训练。当每个线程预先分配有独立的样池时,上述方法可以并行化处理。算法2给出了并行在线增广的详细过程:
为了更好的训练效果,对样本池进行混洗非常重要,但同时它也会减慢网络增广的速度。原因是普通的洗牌方法包含大量随机内存访问,并且不能由CPU缓存加速。如果服务器有多个CPU插槽,速度损失将更加严重。为了缓解这个问题,作者提出了一种伪混洗技术,它以更加缓存友好的方式对相关样本进行洗牌,并显着提高系统的速度。
并行负采样
在嵌入训练阶段,作者将训练任务分为多个子任务并将它们分配给多个GPU设备。子任务必须设计成具有很少的共享数据的形式以最小化GPU之间的同步成本。为了了解如何将模型参数分配给多个GPU而不重叠,作者首先介绍了ε-梯度可交换的概念:
其中,作者将0-梯度可交换称为梯度可交换。由于节点嵌入训练过程的稀疏性质,在网络中会存在很多形成梯度可交换对的集合。例如,对于两个边缘样本集例如,对于两个边缘样本集X1,X2⊆E,如果它们不共享任何源节点或目标节点,则X1和X2是梯度可交换的。即使X1和X2共享一些节点,如果学习率α和迭代次数有界,它们仍然可以是ε-梯度可交换的。
基于节点嵌入中观察到的梯度交换性,作者提出了一种用于嵌入训练阶段的并行负采样算法。对于n个GPU,我们将顶点和上下文行分别划分为n个分区(上图的左上角)。这导致样本池的n×n分区网格,其中每个边缘属于其中一个块。这样,任何不共享行或列的块都是可梯度交换的。只要限制每个块上的迭代次数,同一行或列中的块都是ε-梯度可交换的。
作者将集(episode)定义为并行负抽样中使用的块级步骤。在每一集中,我们分别向n个GPU发送n个正交块及其对应的顶点和上下文分区。然后,每个GPU使用ASGD更新自己的嵌入分区。因为这些块是相互梯度可交换的并且不共享参数矩阵中的任何行,所以多个GPU可以在没有任何同步的情况下同时执行ASGD。在每集结束时,我们从所有GPU收集更新的参数并分配另外n个正交块。这里ε-梯度可交换由n个正交块中的总样本数控制,我们将其定义为集的大小。较小的集有利于ε-梯度可交换嵌入式训练。但是同时也会导致更频繁的同步。因此作者调整了集的大小,以便在速度和ε-梯度可交换之间有一个很好的权衡。上图给出了具有4个分区的并行负采样的示例。
虽然作者使用等于n的分区数量来说明了负采样的有效性,但是并行负采样可以很容易地推广到任何数量大于n的分区的情况,只需在每集中处理n个子组中的正交块。算法3给出了用于多个GPU的混合系统的并行负采样算法的详细步骤:
协作策略
我们的并行负采样使不同的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