PDGNN：通过度数剪枝划分和冗余消除实现高效的微批量图神经网络训练

《ACM Transactions on Architecture and Code Optimization》：PDGNN: Efficient Micro-batch GNN Training via Degree-Pruned Partitioning and Redundancy Elimination

图神经网络的小批量训练方法存在内存与精度权衡问题，本文提出PDGNN系统通过度剪枝优化图分区策略，并采用依赖感知的冗余匹配消除重复计算，在保持模型精度的同时显著降低内存占用，实验显示训练速度提升3.5倍，GPU内存峰值降低46%，CPU内存降低64%。

　　

摘要

图神经网络（GNNs）在基于图的应用中变得至关重要。现有的GNN训练框架主要采用基于采样的小批量方法，通过采样顶点邻域来迭代构建子图进行训练，但这种方法存在内存与准确率之间的权衡。为了解决这个问题，可以应用简单的微批量训练并累积梯度，但这会导致由于重复采样邻居而产生的严重冗余。最近的微批量训练方法通过采用批量级图划分策略来克服这一问题，这些策略将从所有目标顶点采样的全批量子图进行划分以生成微批量。然而，这些方法在平衡划分开销和冗余方面仍然存在困难。

本文提出了PDGNN，这是一种专为单台机器上的高效微批量GNN训练设计的系统，能够减少训练时间和内存消耗。首先，我们提出了一种基于度数剪枝的图划分方法，该方法通过移除高度顶点来高效处理简化和稀疏化的多层二分图，从而显著降低划分开销。其次，我们提出了一种考虑依赖关系的冗余匹配和消除技术，该技术能够最大化微批量之间的冗余匹配，并通过嵌入缓存和重用机制来消除冗余计算，同时保持模型准确性。在真实世界数据集上的实验表明，PDGNN大幅提高了大规模GNN训练的效率，与最先进的批量级训练方法相比，训练时间最多可加快3.5倍，GPU峰值内存消耗降低了46%，CPU峰值内存消耗降低了64%。