已开源!GraphVite 超高速图表示学习系统,1 分钟可学百万节点

2019-08-09 20:48:40 浏览数 (1)

8 月 5 日晚,GraphVite 开发者 @唐建(MILA 实验室助理教授,曾获 ICML 2014最佳论文、WWW16 最佳论文提名) 在社交平台上公布了这个图表示学习系统开源的消息。他表示,在百万节点的图上,使用该系统仅需 1 分钟左右就可以学习节点的表示。该系统的目标是为广泛的嵌入方法系列提供通用和高性能的框架,这将非常有利于图学习算法的研究与部署。雷锋网 AI 开发者将其具体介绍及相关地址编译如下。

GraphVite:高速大规模图形嵌入

GraphVite 是一种通用图形嵌入引擎,专用于各种应用程序中的高速和大规模嵌入式学习。该系统支持多 GPU 并行,它可以扩展到百万级甚至十亿级的图。

GraphVite 可以帮助用户实现:

  • 在统一的平台上重现学习算法
  • 对图形或高维数据进行快速可视化
  • 对大规模图形或知识图谱进行学习
  • 提高原型设计与模型调整的效率

目前,GraphVite 已为 3 种任务提供了完整的训练和评估流程,包括:节点嵌入、知识图谱嵌入和图形&高维数据可视化。

GraphVite 支持的 3 种任务模型

此外,它还涵盖了 9 个主流的模型,以及它们在一系列标准数据集上的基准。通过 Python 界面,用户可以轻松地练习高级图形嵌入算法,并在极短的时间内获得结果。

理解框架

GraphVite 框架由两个部分组成,核心库和 Python wrapper。Python wrapper 可以为核心库中的类提供自动打包功能,并为应用程序和数据集提供了实现。

核心库用 C 11 和 CUDA 实现,并使用 pybind11 绑定到 python 中。它涵盖了 GraphVite 中所有与计算相关类的实现,例如图、求解器和优化器。所有这些成分都可以打包成类,这类似于 Python 接口。

在 C 实现中,Python 有一些不同之处。图和求解器由底层数据类型和嵌入向量长度实现。该设计支持 Python 接口中的动态数据类型,以及对最大化优化编译时(compile-time)。为了方便了对 GraphVite 的进一步开发,开发者还对 C 接口进行了高度抽象。通过连接核心接口,用户可以实现图形的深度学习例程,而无需关注调度细节。

源代码组如下:

  • include/base/*实现基本数据结构
  • include/util/*实现基本用途
  • include/core/*实现优化器、图和求解器的核心接口
  • include/gpu/*实现所有模型的前向和后向传播
  • include/instance/*实现图和求解器的实例
  • include/bind.h 实现 Python 绑定
  • src/graphvite.cu 实例化所有 Python 类

GraphVite 到底有多快?

为了简要介绍 GraphVite 的速度,开发者展示了用 GraphVite 实现的所有模型的基准测试,包括它们的时间和性能。整个测试过程均采用的是具有 24 个 CPU 线程和 4 个 V100 GPU 的服务器。

节点嵌入

开发者在 3 个数据集上进行了节点嵌入模型的实验,这些数据集的规模从百万到 5 亿不等。下图显示了每个数据集的大小,以及嵌入 LINE 模型中所占用时间和 CPU 内存大小。

节点嵌入模型在 3 个数据集上的测试结果

在多标签节点分类的标准任务上,开发者也对学习到的节点嵌入进行了评估。下面展示了每个模型以不同百分比训练数据在 micro-f1 和 macro-f1 两种衡量指标下的值。

Youtube 数据集上的节点嵌入评估

对于大规模数据集的测试结果如下所示。其中,node2vec 模型因内存不足而无法成功运行,因为它需要超过 200 GiB 的内存来为二阶随机游走构建别名表。

Flickr 大规模数据集上的测试结果

Friendster-small 大规模数据集上的测试结果

知识图谱嵌入

对于知识图嵌入,开发者在 4 个标准数据集上对 TransE,DistMult,ComplEx 和 RotatE 进行了基准测试。RotatE 在这些数据集上的训练时间和资源如下图所示。

知识图谱嵌入任务在 4 个标准数据集上测试结果

为了评估知识图谱嵌入,开发者还在链路预测任务上进行了测试,并在测试集上报告每个模型的结果,其中排名指标是根基于过滤结果计算而得。

知识图谱嵌入任务在 FB15k 数据集上评估结果

知识图谱嵌入任务在 FB15k-237 数据集上评估结果

知识图谱嵌入任务在 WN18 数据集上评估结果

图形&高维数据可视化

开发者在两个主流的图像数据集上评估了高维数据可视化的对比结果。LargeVis 所需的训练时间和资源如下表所示。请注意,超过 95%的 GPU 内存占用来自 KNN Graph 的构建,并且可以在必要时减慢速度进行平衡。

高维数据可视化任务在数据集上测试结果

对于 ImageNet,因为它包含 1000 个类,所以开发者根据它们在 WordNet 中的层次结构来可视化类。下面的动画演示了如何在层次结构中遍历 English Setter 类。

在层次结构中遍历 English Setter 类演示

与类似工作的比较

与 GraphVite 相似的系统是 PyTorch-BigGraph,该系统旨在加速大规模知识图嵌入数据。下图是在相同的超参数设置下,对 FB15K 上两个库中实现的模型进行的 Apple-to-Apple 比较。

GraphVite 与 PyTorch-BigGraph系统速度对比

更多 GraphVite 信息: https://graphvite.io/ GitHub 开源地址: https://github.com/DeepGraphLearning/graphvite

0 人点赞