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

混合 CPU-GPU 训练新范式

GPU 虽能高效执行神经网络张量运算，但内存有限。针对超出 GPU 内存容量的大规模图神经网络训练，传统方法需通过 CPU 随机选择节点和边组成小批量数据，并连同节点特征一并传输至 GPU。这种递归邻居采样方式导致大量 CPU-GPU 数据传输，成为模型训练中最耗时的环节。

全局邻居采样(GNS)核心技术

1. 特征缓存机制

2. 在创建小批量数据前，先从全图中采样一组节点并将其特征预加载至 GPU 内存（称为缓存）。生成小批量时优先从缓存中检索邻居节点，仅当缓存不足时才从 CPU 补充获取。

3. 高权重节点优先策略

4. 采样概率与节点度数成正比，既确保高频连接节点优先缓存，又保留部分低频节点。通过数学证明，该方法能收敛至与完全随机采样相同的模型最优性能。

5. 子图索引优化

6. 在 CPU 端构建包含所有缓存节点及其直连邻居的子图，通过快速查询替代昂贵的集合运算，显著提升邻居匹配效率。

实验验证

• 在 5 个数据集上的测试表明，GNS 在混合 CPU-GPU 环境下比次优方案快 2 倍以上

• 两种传统采样方法需耗费 10 倍时间才能达到相近精度

• 当前支持单 CPU-GPU 架构，多 GPU 分布式训练的扩展研究正在进行中

性能突破关键

通过减少 90%以上的 CPU-GPU 数据传输，GNS 将训练瓶颈从数据搬运转向实际计算，为十亿级节点规模的图神经网络训练提供实用解决方案。该方法已发表于 KDD 2021 会议，相关代码实现同步公开。更多精彩内容 请关注我的个人公众号 公众号（办公 AI 智能小助手）公众号二维码

办公AI智能小助手