北大、微软亚洲研究院：高效的大规模图神经网络计算

新智元报道

来源：arXiv

编辑：肖琴

【新智元导读】第一个支持图神经网络的并行处理框架出现了！北京大学、微软亚洲研究院的研究人员近日发表论文，提出NGra，这是第一个支持大规模GNN的系统。

GNN（图神经网络）代表了一种新兴的计算模型，这自然地产生了对在大型graph上应用神经网络模型的需求。

但是，由于GNN固有的复杂性，这些模型超出了现有深度学习框架的设计范围。此外，这些模型不容易在并行硬件（如GPU）上有效地加速。

近日，北京大学、微软亚洲研究院的多位研究人员在arXiv上发布了一篇新论文，提出了解决这些问题的有效方案。

论文题为Towards Efficient Large-Scale Graph Neural Network Computing：

论文地址：https://arxiv.org/pdf/1810.08403.pdf

作者表示：“我们提出NGra，这是第一个基于图形的深度神经网络并行处理框架。”

NGra描述了一种新的SAGA-NN模型，用于将深度神经网络表示为顶点程序（vertex programs） ，其中每一层都在明确定义的图形操作阶段（Scatter，ApplyEdge，Gather，ApplyVertex）。

这个模型不仅允许直观地表示GNN，而且还可以方便地映射到高效的数据流表示。NGra通过GPU核心或多GPU的自动图分区和基于chunk的流处理透明地解决了可扩展性挑战，仔细考虑了数据局部性、数据移动以及并行处理和数据移动的重叠。

NGra通过在GPU上进行高度优化的Scatter / Gather操作进一步提高了效率，尽管它具有稀疏性。我们的评估表明，NGra可以扩展到现有框架无法直接处理的大型实际图形，而在TensorFlow的multiple-baseline设计上，即使在小规模上也可以实现约4倍的加速。

第一个支持大规模GNN的系统

NGra是第一个支持大规模GNN（图神经网络）的系统，这是一个在GPU上可扩展、高效的并行处理引擎。

NGra自然地将数据流（dataflow）与顶点程序抽象（vertex-program abstraction）结合在一个新模型中，我们将其命名为SAGA-NN（Scatter-ApplyEdge-Gather-ApplyVertex with Neural Networks）。

虽然SAGA可以被认为是GAS（Gather-Apply-Scatter）模型的变体，但SAGA-NN模型中的用户定义函数允许用户通过使用数据流抽象来表示对vertex或edge数据（被视为tensors）的神经网络计算，而不是专为传统图形处理而设计（例如PageRank、 connected component和最短路径等算法）

与DNN一样，高效地使用GPU对于GNN的性能至关重要，而且由于要处理的是大型图形结构，这一点更为重要。为了实现超出GPU物理限制的可扩展性，NGra将图形（顶点和边缘数据）透明地划分为块（chunk），并将SAGA-NN模型中表示的GNN算法转换为具有chunk粒度的运算符的dataflow graph，从而在单个GPU或多个GPU上启用基于chunk的并行流处理。

NGra engine的效率在很大程度上取决于NGra如何管理和调度并行流处理，以及在GPU上关键图形传播运算符Scatter和Gather的实现。

NGra非常注重数据局部性，以最大限度地减少GPU内存中的数据交换，并在GPU内存中最大化数据块的重用，同时将数据移动和计算以流的方式重叠。

对于多GPU的情况，它使用 ring-based streaming机制，通过直接在GPU之间交换数据块来避免主机内存中的冗余数据移动。

与其他基于GPU的图形引擎关注的传统图形处理场景不同，在GNN场景中，可变顶点数据本身可能无法容纳到GPU设备内存中，因为每个顶点的数据可以是特征向量（ feature vector）而不是简单的标量（scalar）。因此，我们的方案更倾向于在每个顶点数据访问中利用并行性，从而提高内存访问效率。

我们通过使用vertex-program abstraction和图形传播过程的自定义运算符扩展TensorFlow，从而实现NGra。

我们利用单个服务器的主机内存和GPU的计算能力，证明NGra可以扩展以支持大型图形的各种GNN算法，其中许多是现有深度学习框架无法直接实现的。

与小型graph上的TensorFlow相比，它可以支持GPU，NGra可以获得最多4倍的加速。我们还广泛评估了NGra的多重优化所带来的改进，以证明其有效性。

接下来的部分将描述SAGA-NN编程抽象，NGra系统的组件，以及NGra的实现和评估。

NGra程序抽象

基于图（graph）的神经网络（GNN）是根据图形结构定义的一类通用神经网络架构。

图中的每个顶点或边可以与张量数据（通常是vector）相关联，作为其特征或嵌入。GNN可以堆叠在多个层中，迭代传播过程在同一个图上逐层进行。

在图的每个层中，顶点或边缘要素沿边缘变换和传播，并在目标顶点聚合，以生成下一层的新要素。转换可以是任意的DNN计算。

图还可以包含每个顶点，每个边缘或整个图形的标签，用于计算顶层的损失函数。然后从底层到顶层执行前馈计算（feedforward computation）和反向传播。

图1描述了一个2层的GNN的前馈计算。

图1

我们使用Gated Graph ConvNet（G-GCN）算法作为一个具体示例。 Graph ConvNet概括了卷积运算的概念，通常应用于图像数据集，用于处理任意图形（例如knowledge graph）。Gated Graph ConvNet进一步结合了门控机制，因此模型可以了解哪些边对学习目标更重要。

G-GCN每一层的前馈计算如图2所示：

图2：SAGA-NN模型中，Gated Graph ConvNet的layer

，其中⊗指矩阵乘法。

NGra系统的组成

NGra提供dataflow和vertex program abstractions的组合作为用户界面。

NGra主要包括：

1. 一个前端，它将SAGA-NN模型中实现的算法转换为块粒度数据流图（chunk-granularity dataflow graph），使GPU中大型图的GNN计算成为可能；
2. 一个优化层，它产生用于最小化主机和GPU设备存储器之间的数据移动的调度策略，并识别融合操作和删除冗余计算；
3. 一组有效的传播操作内核，支持基于流的处理，以将GPU中的数据移动和计算重叠；
4. dataflow execution runtime。NGra主要利用现有的基于数据流的深度学习框架来处理dataflow execution runtime。

图3：SAGA-NN Stages for each layer of GN

NGra的优化

图4描述了ApplyEdge阶段中矩阵乘法运算：

图4

图5显示了优化的dataflow graph，其中矩阵乘法移入ApplyVertex stage：

图

图7是多GPU的架构

图7：多GPU架构

NGra的评估

我们在TensorFlow (v1.7) 上实现NGra，使用大约2,900行C++代码和3000行Python代码。NGra通过前端扩展TensorFlow，将SAGA-NN程序转换为chunk-granularity dataflow graph，几个scatter/gather 运算符，以实现高效的图传播，以及ring-based的流调度方案。

以下是评估结果。评估证明了NGra的高效和可扩展性，以及与state-of-the-art的系统TensorFlow的比较。

表1：数据集 (K: thousand, M: million)

图13：TensorFlow(TF)，cuSPARSE和NGra(NG) 在不同密度graphs上的传播内核时间

表2：与TensorFlow的迭代时间比较（ms）

图14：不同应用程序的Streaming scheduling策略比较。(Data: reddit middle)

图15：NGra在不同应用程序的扩展性能

图16：在大型图上使用不同应用程序加速NGra

结论

GNN代表了一种新兴的计算模型，这自然地产生了对在大型graph上应用神经网络模型的需求。由于GNN训练固有的复杂性，支持高效的、可扩展的并行计算是很困难的。

NGra是第一个支持GNN的并行处理框架，它使用新的编程抽象，然后将其映射和优化为数据流，进而在GPU上高效执行。

论文地址：https://arxiv.org/pdf/1810.08403.pdf