基于采样的图神经网络数据加载系统

作     者：熊颖彤 

作者单位：华东师范大学 

学位级别：硕士

导师姓名：翁楚良

授予年度：2023年

学科分类：08[工学] 081104[工学-模式识别与智能系统] 0811[工学-控制科学与工程] 

主      题：图神经网络 数据加载 基于采样的图神经网络训练 

摘      要：随着图神经网络受到的关注越来越多,各种图神经网络模型层出不穷。事实证明,图神经网络已经成功地被应用于各个领域中,例如推荐系统、社交网络分析、蛋白质分子分析等。然而,随着图数据尺寸爆炸性的增长,传统的图神经网络训练由于邻域爆炸问题的存在,面临着计算和存储两方面的压力。同时,基于全图的训练方式使得模型参数在一次完整迭代中只更新一次,导致模型的收敛速度可能受到影响。因此,基于采样的图神经网络训练被提出来,以缓解图神经网络训练过程中所面临的计算和存储压力并提升模型的收敛速度。为了提升基于采样的图神经网络的训练效率,不少研究工作提出了多种采样算法,试图从算法的角度提升训练效率。这些研究工作提出的算法理论上能为图神经网络训练带来性能提升,却忽略了算法设计和系统实现之间存在的差距,使得在实际训练过程中依然存在着性能瓶颈。本文针对基于采样的图神经网络训练进行优化,以提升网络训练的效率。首先,本文通过实验分析得出目前制约基于采样的图神经网络训练的两个瓶颈为采样和数据加载。接下来,本文针对这两个瓶颈设计了相应的优化方法。实验结果表明,本文提出的优化方法能够提升端到端的图神经网络训练效率。本文的主要工作如下:(1)基于采样的图神经网络训练定量分析:为了找出基于采样的训练方式中存在的瓶颈,本文在当前主流的图神经网络系统上进行测试,分别对采样和数据加载进行定量分析,识别出采样和数据加载是训练过程中存在的主要瓶颈。另一方面,本文发现在基于采样的训练方式中,图数据存在着可待利用的数据局部性。基于以上分析,本文分别对数据采样和数据加载进行了优化。(2)基于GPU的高效邻居节点采样器:为了提升图采样的性能,本文将采样过程从CPU端迁移到GPU端,并利用GPU的高度并行性设计了一个高效的采样器以实现快速高效的采样。同时,为了进一步提升采样效率,本文结合图数据特征和硬件特征,设计了一个针对图拓扑数据的存储策略。实验显示,相比于其他图神经网络系统,GPU采样器能够提升图数据采样性能。(3)一个高效的数据加载器:本文设计了一个高效的数据加载器以优化数据加载。结合上文分析得到的可待利用的数据局部性,本文将mini-batch中的节点划分为共享节点和非共享节点。对于共享节点,本文通过控制GPU中特征数据的生命周期的方式,重复利用共享节点的特征数据,从而减少CPU和GPU之间的冗余数据传输;对于非共享节点,本文设计了一个高效的特征聚合器来提升特征聚合的性能。实验显示,相比于逐mini-batch的数据传输,本文利用图数据的数据局部性能够提升CPU和GPU之间的数据传输吞吐。(4)数据加载优化系统的实现:本文将上述优化工作实现为一个高效的数据加载优化系统,Fast Data Loader(FDL)。FDL提供了简单易用的前端接口,使得用户在使用过程中尽可能少地减少对相应代码的修改。实验显示,FDL能够有效的提升基于采样的图神经网络训练效率。综上所述,本文首先分析出基于采样的图神经网络训练中的瓶颈为采样和数据加载以及图数据在mini-batch之间体现出的数据局部性。结合上述分析,本文从数据采样和数据加载这两个角度进行优化,最终提高了端到端的图神经网络训练效率。