[论文解读] Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks

Minjie Wang, Da Zheng+13 more|arXiv (Cornell University)|Sep 3, 2019

Advanced Graph Neural Networks参考文献 52被引用 713

一句话总结

DGL 提供一个面向图的 GNN 框架，将计算抽象为通用的 g-SpMM 和 g-SDDMM 基元，从而实现框架中立的部署（PyTorch、TensorFlow、MXNet），并在基准测试中提供更高的速度和更低的内存开销。

ABSTRACT

Advancing research in the emerging field of deep graph learning requires new tools to support tensor computation over graphs. In this paper, we present the design principles and implementation of Deep Graph Library (DGL). DGL distills the computational patterns of GNNs into a few generalized sparse tensor operations suitable for extensive parallelization. By advocating graph as the central programming abstraction, DGL can perform optimizations transparently. By cautiously adopting a framework-neutral design, DGL allows users to easily port and leverage the existing components across multiple deep learning frameworks. Our evaluation shows that DGL significantly outperforms other popular GNN-oriented frameworks in both speed and memory consumption over a variety of benchmarks and has little overhead for small scale workloads.

研究动机与目标

• 将 GNN 计算萃取为一小组通用、可优化的基元（g-SpMM 和 g-SDDMM）。
• 将图作为中心编程抽象，以简化用户代码并实现透明优化。
• 实现框架中立，以便在 PyTorch、TensorFlow 和 MXNet 之间移植，同时保持性能。
• 通过优化的并行化策略和融合计算实现高速度和内存效率。

提出的方法

• 将 GNN 的信息传递形式化为广义 SpMM（g-SpMM）和广义 SDDMM（g-SDDMM）操作。
• 开发并行化策略（针对 g-SpMM 的节点并行、针对 g-SDDMM 的边并行）并讨论格式（CSR/CSC/ COO）以提升性能。
• 实现 DGLGraph 作为中心数据结构，配备框架 shim 并自动切换格式以优化前向和反向传播。
• 通过最小化框架变更实现框架中立设计；使用 DLPack 进行张量共享并创建可微分算子以进行反向传播。
• 暴露面向图的 API（g.update_all、g.apply_edges）以在不实现场间张量的情况下，将信息传递与用户定义的函数组合起来。

实验结果

研究问题

• RQ1 图神经网络计算是否能够被两个通用基元（g-SpMM 和 g-SDDMM）有效捕获？
• RQ2 以图为中心、框架中立的包在多样的 GNN 模型和数据集上是否提供更高的速度和内存效率？
• RQ3 哪些并行化策略最适合利用 GPU/CPU 硬件实现 g-SpMM 和 g-SDDMM？
• RQ4 在一个框架中立的库内，将 GNN 模型在 PyTorch、TensorFlow 和 MXNet 之间迁移需要多少移植工作？

主要发现

• DGL 在基准测试中的速度和内存消耗方面显著优于其他流行的面向 GNN 的框架。
• g-SpMM 和 g-SDDMM 内核融合计算和聚合，减少内存传输，使在比 PyG 能处理的规模更大的图上进行训练成为可能（如 ML-10m，PyG 内存不足）。
• 在 CPU 上，DGL 相对于 PyG 实现显著加速（1.9x 到 64x 的提升），原因是其内核对 CPU 的利用率更高。
• 使用 CS（cluster sampling）的小批量训练在 DGL 上对 GAT 能带来约 1.56x 的加速。
• DGL 提供低开销的框架中立后端，相对于特定于框架的实现具有竞争力的性能。