GNNLab：基于空间共享思想设计的新型 GNN 系统

思想设计的新型 GNN 系统"/>

GNNLab：基于空间共享思想设计的新型 GNN 系统

本文介绍 GraphScope 团队与上海交通大学陈榕教授团队的学术合作项目——基于 GPU 的图神经网络采样训练系 GNNLab。该成果论文 GNNLab: A Factored System for Sample-based GNN Training over GPUs 发表于计算机系统领域顶级学术会议 EuroSys’2022。

背景

近年来，以 GCN、GraphSage 和 PinSage 为代表的图神经网络（Graph Neural Networks， GNNs）在推荐系统、社交网络信息挖掘、蛋白质功能分析等领域取得了巨大的成功，各种各样的 GNN 训练系统（例如 DGL 和 PyG）如雨后春笋般不断涌现。GNN 将深度学习技术与图（Graph）数据进行结合，与传统的深度学习不同，图中的顶点需要不断地聚合图中邻居顶点的特征。下图简要展示了 GNN 与传统的深度学习模型训练流程的区别。

对于一个 GNN 模型来说，其输入的图数据包含图结构数据（graph topological data）和图特征数据（graph feature data）两部分。随着图数据规模的不断增大，如何高效地训练 GNN 成为一个难题，目前在大规模的数据集上普遍采用基于图采样的 mini-bacth 训练方法。该方法包含 Sample、Extract 和 Train 三个阶段，也被称为 SET 模型。此方法将所有的训练节点划分为若干个小的 mini-bacth，对于每个 mini-batch，我们首先在 Sample 阶段利用图采样算法为每个顶点选取若干个邻居构成一个 sample，然后在 Extract 阶段在为 sample 中的每个顶点抽取对应的特征后进入 Train 阶段，即进行模型训练。整个训练流程不断重复上述过程，直到模型的准确率达到期望