美团图神经网络框架Tulong，图计算库MTGraph

图神经网络训练框架的实践和探索

1. 前言

深度学习技术将图像、文本、语音等多种多样的数据转化为稠密的向量表示，提供了表示数据的另一种方式。

由于图数据特有的稀疏性（图的所有节点对之间只有少量边相连），直接使用通用的深度学习框架（例如TensorFlow和PyTorch）训练往往性能不佳。

针对图神经网络的深度学习框架应运而出：PyG (PyTorch Geometric)[6]和DGL (Deep Graph Library)[7]等开源框架大幅提升了图神经网络的训练速度，并且降低了资源消耗[17][18]，拥有活跃的社区支持。

1.1 一个“好用”的图神经网络框架

至少具备以下特点。

（1）完善支持当前流行的图神经网络模型。

（2）以合理的代价支持大规模图上的模型训练。

希望单机即可在合理的时间内训练百亿边规模的模型

（3）与业务系统无缝对接。

图神经网络的完整落地流程至少包括：基于业务数据构图、离线训练和评测模型、线上推理、业务指标观测等步骤。合适的工具能提升对接业务数据的效率，然而现有的图神经网络框架大多聚焦在模型的离线训练和评测，缺乏此类工具。

（4）研发人员易于上手，同时提供充足的可扩展性。

通过简单地配置即能完成多数任务。在此基础上，对于一些特殊的建模需求，也能提供适当的支持。

1.2 美团的解决方案

图神经网络框架Tulong以及配套的图学习平台

对当前流行的图神经网络模型进行了细粒度的剖析，归纳总结出了一系列子操作，实现了一套通用的模型框架。简单修改配置即可实现许多现有的图神经网络模型。
高度可配置化的训练和评测，从参数初始化到学习率，从模型结构到损失函数类型…

2. 系统概览

（1）图以及深度学习引擎

我们把图神经网络的底层算子分为三类：图结构查询、稀疏张量计算和稠密张量计算。

我们开发了图计算库MTGraph提供图数据的存储和查询功能，深度优化了内存占用和子图采样速度。

MTGraph兼容PyTorch和DGL，用户可以在MTGraph的基础上直接编写基于DGL的模型代码。

（2）Tulong框架

Tulong框架首先封装实现了训练图神经网络所需的基本组件，包括图和特征数据的预处理流程、子图采样器、通用的GNN模型框架，以及包括训练和评测在内的基础任务。

基于上述组件，Tulong框架提供丰富的预定义模型和训练/推理流程，用户通过修改配置文件即可在业务数据上训练和评测GNN模型。

3. 模型框架

我们从工程实现的角度，归纳总结了当前主流图神经网络模型的基本范式，实现一套通用框架，以期涵盖多种GNN模型。

3.1 同质图

同质图(Homogeneous Graph)可以定义为节点集合和边集合：**G=(V,E)**，一条边(u,v)∈E表示节点u与节点v相连。节点和边上往往还附加有特征，我们记x(u,v)为边(u,v)的特征。

前向计算过程的计算范式：

上述计算范式仍然分为生成消息、聚合消息、更新当前节点三个步骤，具体包括：

层次维度的聚合函数ρL(⋅)：用于聚合同一节点在模型不同层次的表示。例如，多数GNN模型中，层次维度的聚合函数为上一层的节点表示；而在JKNet[10]中，层次维度的聚合函数可以设定为LSTM[11]。
消息函数ϕ(⋅)：结合起始节点和目标节点，以及边的特征，生成用于消息传递的消息向量。
节点维度的聚合函数ρN(⋅)：汇集了来自邻居节点N(v)的所有消息向量。值得注意的是，N(v)也可以有不同的实现。例如，在GCN中为所有邻居节点，而在GraphSage[9]中为邻居节点的子集。
更新函数ψ(⋅)：用于聚合节点自身在上一层和当前层的表示。

不难看出，上述计算范式可以覆盖当前大多数GNN模型。在工程实践中，我们将上述函数进一步分拆细化，预先提供了多种高效的实现。通过配置选项即可实现不同的组合搭配，从而实现多数主流的GNN模型。

3.2 异质图

相比于同质图，异质图(Heterogeneous Graph)扩充了节点类型和边类型。

比如，学术引用网络[13]中包含论文、作者、机构等类型的节点，节点直接通过“论文引用其他论文”、“作者撰写论文”、“作者属于机构”等类型的边相连，如下图2所示：

图2 同质图与异质图的比较

我们把异质图视为多个二分图的叠加，每一个二分图对应于一种边类型。上述的学术引用网络可以表示成“论文-引用-论文”、“作者-撰写-论文”、“作者-属于-机构”，共计三个二分图，同质图的GNN模型框架稍加修改即可在二分图上应用。

在此基础上，一个节点在不同的二分图中会产生不同的表示。我们进一步提出边类型维度的聚合函数ρR(⋅)，用于聚合节点在不同二分图中的表示（如下图3所示）。框架中同样提供边类型维度聚合函数的多种实现，可以通过配置选项调用。例如，要实现RGCN，可以在二分图上应用GCN，然后在边类型维度上取平均。

图3 异质图模型框架

3.3 动态图

动态图(Dynamic Graph)是指随时间变化的图。与之相对的，上述的同质图和异质图可以称为静态图。动态图上的GNN模型旨在生成给定时间下的节点表示H(t)。根据时间粒度的粗细，动态图可分为离散时间动态图和连续时间动态图。

在离散时间动态图中，时间被划分为多个时间片（例如以天/小时划分），每个时间片对应一个静态的图。离散时间动态图的GNN模型通常在每个时间片上单独应用GNN模型，然后聚合节点在不同时间的表征[14]。我们把聚合过程抽象为离散时间维度的聚合函数ρT(⋅)，同样提供预定义的实现。

此外，Tulong框架还提供离散时间动态图数据的加载和管理机制，仅在内存中保留必须的时间片，降低硬件资源的消耗。

图4 离散时间动态图GNN模型框架

在连续时间动态图中，每条边附有时间戳，表示交互事件发生的时刻。相比于静态图，连续时间动态图中的消息函数ϕ(⋅,t,et)还依赖于给定样本的时间戳以及边的时间戳。此外，邻居节点N(v,t)必须与时间有关，例如邻居节点中不能出现t时刻之后才出现的节点。

针对此问题，我们开发了多种连续时间动态图上的邻居节点采样器，可以在指定的时间范围内，高效地采样邻居节点。

图5 连续时间动态图GNN模型框架

以上分析了同质图、异质图和动态图的计算范式，我们从中抽取出通用的函数（算子），包括消息函数、聚合函数、更新函数、邻居节点函数，并给出多种预定义的实现。框架用户通过配置选项即可拼装组合算子，从而实现需要的GNN模型。

4. 训练流程框架

训练GNN模型通常包括加载数据、定义GNN模型、训练和评测、导出模型等流程。

框架被分为两层：基础组件和流程组件。

基础组件聚焦于单一的功能，例如图数据组件只维护内存中的图数据结构，不提供图上的采样或张量计算功能；图上的采样功能通过图采样器来提供。

流程组件通过组装基础组件提供较为完整的数据预处理、训练和评测流程，例如训练流程组合了图数据、图采样器、GNN模型等组件，提供完整的训练功能。

图6 训练流程框架

更上一层，我们提供多种流程配置模板和GNN模型模板。模板对外暴露若干超参，例如训练数据路径、模型类型、学习率等参数，结合用户指定的超参后就可以完整定义一次训练任务。换言之，基于模板和参数即可完整复现一次GNN模型实验。框架将会解析这些配置，并生成可执行的应用。

举例来说，用户可以选择GraphSage模型的配置模板，以及链接预测任务的训练模板，指定模型层数和维度，以及训练评测数据路径，即可开始训练基于GraphSage的链接预测模型。

5. 性能优化

高效训练数十亿乃至百亿边规模的GNN模型

5.1 图数据结构优化

我们设计实现了更为紧凑的图数据结构，提升了单机可承载的图规模。

我们借助图压缩技术降低内存占用。不同于常规的图压缩问题，GNN的场景下需要支持随机查询操作。例如，查询给定节点的邻居节点；判断给定的两个节点在图中是否相连。

我们对此提出的解决方案包括两部分：

图数据预处理和压缩：首先分析图的统计特征，以轻量级的方式对节点进行聚类和重新编号，以期让编号接近的节点在领域结构上也更为相似。随后调整边的顺序，对边数据进行分块和编码，产生“节点-分块索引-邻接边”层次的图数据文件（如下图7所示）。最后，如果数据包含节点特征或边特征，还需要将特征与压缩后的图对齐。

图7 压缩后的图数据结构

图的随机查询：查询操作分为两步：首先定位所需的边数据块，然后在内存中解压数据块，读取所查询的数据。例如在查询节点u和v是否相连时，首先根据两个节点的编号计算边数据块的地址，解压数据块后获得少量候选邻接边（通常不多于16条），然后查找是否包含边(u,v)。

5.2 子图采样优化

子图采样是GNN模型训练的性能瓶颈之一。

我们分别针对静态图和动态图，设计实现了多种高效的邻居节点采样算法。主要的优化手段包括：

随机数发生器：相比于通信加密等应用，图上的采样对于随机数发生器的“随机性”并没有苛刻的要求。我们适当放松了对随机性的要求，设计实现了更快速的随机数发生器，可以直接应用在有放回和无放回的采样操作中。
概率量化：有权重的采样中，在可接受的精度损失下，将浮点数表示的概率值量化为更为紧凑的整型。不仅降低了采样器的内存消耗，也可以将部分浮点数操作转化为整型操作。
时间戳索引：动态图的子图采样操作要求限定边的时间范围。采样器首先对边上的时间戳构建索引，采样时先根据索引确定可采样边的范围，然后再执行实际的采样操作。

6. 图学习平台

数据集管理：从业务数据构造图是模型开发的第一步，图学习平台提供基于Spark的构图功能，可以将Hive中存储的业务数据转化为Tulong自定义的图数据格式。业务数据经常以事件日志的方式存储，如何从中抽象出图，有大量的选择。例如，在推荐场景中，业务日志包含用户对商家的点击和下单记录，除了把”用户-点击-商家”的事件刻画为图以外，还可以考虑刻画短时间内共同点击商家的关系。除此之外，还可以引入额外的数据，比如商家的地理位置、商家在售的菜品等。究竟使用何种构图方案，需要经过实验才能确定。对此，图学习平台提供了图形化的构图工具（如下图10所示），帮助用户梳理构图方案；同时还提供图数据集的版本管理，方便比较不同构图方案的效果。

实验管理：确定图数据之后，建模方案和训练策略是影响最终效果的关键。例如，应该用何种GNN模型？损失函数如何选取？模型超参和训练超参如何确定？这些问题也需要经过大量实验才能回答。基于Tulong框架，建模方案和训练策略可以通过一组配置来控制。图学习平台提供配置的可视化编辑器和版本管理功能，方便比较不同的方案的优劣。

流程管理：有了图数据集和建模/训练方案后，还需要让整个流程自动化。这是模型上线的必要条件，同时也有利于团队成员复现彼此的方案。图学习平台针对常见的“构图、训练、评测、导出”流程提供了自动化的调度，在适当的时候可以复用前一阶段的结果，以提升效率。例如，如果数据集的定义没有变化，可以跳过Spark构图阶段直接使用已有的图数据。此外，针对模型上线的需求，平台提供构图和建模方案整合和定时调度等功能。

9. 参考文献

[1] Cai, Hongyun, Vincent W. Zheng, and Kevin Chen-Chuan Chang. “A comprehensive survey of graph embedding: Problems, techniques, and applications.” IEEE Transactions on Knowledge and Data Engineering 30, no. 9 (2018): 1616-1637.
[2] Perozzi, Bryan, Rami Al-Rfou, and Steven Skiena. “Deepwalk: Online learning of social representations.” In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 701-710. 2014.
[3] Grover, Aditya, and Jure Leskovec. “Node2vec: Scalable feature learning for networks.” In Proceedings of the 22nd ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 855-864. 2016.
[4] Kipf, Thomas N., and Max Welling. “Semi-supervised classification with graph convolutional networks.” International Conference on Learning Representations (2017).
[5] Wu, Zonghan, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S. Yu Philip. “A comprehensive survey on graph neural networks.” IEEE transactions on neural networks and learning systems 32, no. 1 (2020): 4-24.
[6] https://github.com/pyg-team/pytorch_geometric
[7] https://www.dgl.ai/
[8] Chen, Jie, Tengfei Ma, and Cao Xiao. “FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling.” In International Conference on Learning Representations (2018).
[9] Hamilton, Will, Zhitao Ying, and Jure Leskovec. “Inductive representation learning on large graphs.” Advances in neural information processing systems 30 (2017).
[10] Xu, Keyulu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. “Representation learning on graphs with jumping knowledge networks.” In International Conference on Machine Learning, pp. 5453-5462. PMLR, 2018.
[11] Hochreiter, Sepp, and Jürgen Schmidhuber. “Long short-term memory.” Neural computation 9, no. 8 (1997): 1735-1780.
[12] https://github.com/snap-stanford/GraphGym
[13] https://ogb.stanford.edu/
[14] Sankar, Aravind, Yanhong Wu, Liang Gou, Wei Zhang, and Hao Yang. “Dysat: Deep neural representation learning on dynamic graphs via self-attention networks.” In Proceedings of the 13th International Conference on Web Search and Data Mining, pp. 519-527. 2020.
[15] Xu, Da, Chuanwei Ruan, Evren Korpeoglu, Sushant Kumar, and Kannan Achan. “Inductive representation learning on temporal graphs.” International Conference on Learning Representations (2020).
[16] https://github.com/dmlc/dgl/tree/master/dglgo
[17] Wang, Minjie, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou et al. “Deep graph library: A graph-centric, highly-performant package for graph neural networks.” arXiv preprint arXiv:1909.01315 (2019).
[18] Fey, M. and Lenssen, J. E. “Fast graph representation learning with PyTorch Geometric.” In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
[19] Schlichtkrull, Michael, Thomas N. Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. “Modeling relational data with graph convolutional networks.” In European semantic web conference, pp. 593-607. Springer, Cham, 2018.

GNN训练框架入门 美团图神经网络框架Tulong