大规模图神经网络的高效训练新方法

在近期于知识发现与数据挖掘会议（KDD）上发表的一篇论文中，作者介绍了一种用于训练图神经网络模型的新采样策略，该策略结合了CPU和GPU进行运算。在这种现实应用中常见的环境下，新方法减少了从CPU传输到GPU的数据量，从而大幅提升了效率。实验中，根据所使用的数据集不同，该方法的速度比以往方法快2到14倍，同时达到了同样高甚至更高的准确率。

混合CPU-GPU训练
GPU是执行神经网络训练所用张量操作的最高效方式，但其内存有限。为了在那些大到无法完全放入GPU内存的图上训练图神经网络，通常采用CPU来创建由随机选取的图节点和边组成的小批量数据，然后将这些数据连同描述每个节点的数据（即节点特征）一起发送给GPU。
为了生成一个小批量数据，需要为每个目标节点（即正在被嵌入的节点）采样其邻居节点，并且在必要时，还需为这些被采样的邻居节点进一步采样它们的邻居。这种递归式的邻居采样所生成的小批量数据，需要CPU和GPU之间进行大量的数据传输。论文中报告的一系列实验表明，采用现有的采样策略时，将小批量数据的节点特征从CPU复制到GPU是模型训练中最耗时的环节。

全局邻居采样
这种新的采样方法被称为GNS，即全局邻居采样。它通过大幅减少从CPU传输到GPU的数据量来提升效率。
其基本策略是，在创建小批量数据之前，先从整个图中采样一组节点，并将它们的特征加载到GPU内存中；这组节点特征被称为“缓存”。当创建小批量数据时，通过直接从缓存中检索已存在的邻居节点来完成邻居采样。只有当缓存中包含的邻居节点数量不足时，才会从CPU中获取额外的节点。
为了提高相关邻居节点在缓存中被找到的概率，该方法会优先采样度数高（即与其他节点连接数量多）的节点。采样概率与节点度数成正比，因此缓存中也会包含一定数量度数相对较低的节点。
由于已知缓存中节点的采样概率，在嵌入过程中可以对缓存中的节点进行加权，以确保能够很好地近似出考虑所有邻居节点后所能得到的嵌入结果。
论文中证明，这种方法在收敛到最优模型性能的效率上与使用真正随机采样的方法相当。这意味着，对高度数节点的偏向以及为多个小批量重复使用同一缓存的做法，都不会损害模型的性能。
一个重要的考量是如何高效地识别出与特定小批量数据相关的缓存节点。一种潜在的方法是计算给定节点的邻居列表与缓存中节点列表的交集。然而，这种计算开销很大。作为替代方案，在CPU上创建了一个子图，该子图包含缓存中的所有节点以及它们所有的直接邻居。当组装一个小批量数据时，只需从这个子图中查找每个节点的已缓存邻居即可。

在实验中，将该采样策略与其他三种方法在五个数据集上进行了比较，发现在混合CPU-GPU环境下，该方法在所有数据集上的速度至少是次优策略的两倍。在训练以达到可比的准确率时，另外三种采样策略中有两种的速度始终比新方法慢一个数量级。
在实验中，仅限于使用单个CPU和单个GPU。在正在进行的工作中，正在考虑如何将该方法推广到多GPU和分布式训练。例如，是否可以在不同的GPU上缓存不同的节点集，并高效地将每个小批量数据定向到其缓存匹配度最佳的GPU？FINISHED
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）