科研之友
微信
新浪微博
Facebook
分享链接
摘要
图神经网络(graph neural network, GNN)是处理图数据的重要方法.由于计算复杂、图数据容量大,在大规模图上训练图神经网络依赖于CPU-GPU协作和图采样训练方法,其中图结构和特征数据存储在CPU内存中,而采样得到的子图及其特征则传输至GPU进行训练.然而,这种方法面临着严重的图特征数据加载瓶颈,显著降低了端到端训练性能,且图特征占用过多内存,严重限制了可训练的图规模.为了解决这些问题,我们提出了基于输入特征稀疏化的数据加载方法,显著减少CPU内存占用和跨PCIe总线传输的数据量,大幅缩短数据加载时间,加速GNN的训练,使其可以充分利用GPU计算资源.针对图特征和GNN计算特性,我们提出了适用于图特征数据的稀疏化方法,在压缩比和模型准确度之间达到平衡.我们在3个常见GNN模型和3个不同规模的数据集上进行了实验评估,包括最大的公开数据集之一MAG240M.结果表明,此方法将特征尺寸减小了一个数量级以上,并实现1.6–6.7倍的端到端训练加速,而模型准确度的降低不超过1%.此外,在仅使用4个GPU的情况下,仅需40 min就可以在MAG240M上完成GraphSAGE模型的训练并达到目标准确度.
