[Paper Review] 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 Networks52 references713 citations

TL;DR

DGL presents a graph-centric framework for GNNs that abstracts computations into general g-SpMM and g-SDDMM primitives, enabling framework-neutral deployment (PyTorch, TensorFlow, MXNet) and delivering superior speed and memory efficiency across benchmarks.

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.

Motivation & Objective

• Distill GNN computations into a small set of generic, optimizable primitives (g-SpMM and g-SDDMM).
• Make graph the central programming abstraction to simplify user code and enable transparent optimization.
• Design for framework neutrality to ease porting across PyTorch, TensorFlow, and MXNet while preserving performance.
• Achieve high speed and memory efficiency through optimized parallelization strategies and fused computations.

Proposed method

• Formalize GNN message passing as generalized SpMM (g-SpMM) and generalized SDDMM (g-SDDMM) operations.
• Develop parallelization strategies (node parallel for g-SpMM, edge parallel for g-SDDMM) and discuss formats (CSR/CSC/ COO) for performance.
• Implement DGLGraph as the central data structure with framework shims and automatic format switching to optimize forward and backward passes.
• Provide framework-neutral design by minimizing framework changes; use DLPack for tensor sharing and create differentiable operators for backpropagation.
• Expose graph-centric APIs (g.update_all, g.apply_edges) to compose message-passing with user-defined functions without materializing intermediate tensors.

Experimental results

Research questions

• RQ1 Can graph neural network computations be effectively captured by two generalized primitives (g-SpMM and g-SDDMM)?
• RQ2 Does a graph-centric, framework-neutral package deliver superior speed and memory efficiency across a diverse set of GNN models and datasets?
• RQ3 What parallelization strategies best exploit GPU/CPU hardware for g-SpMM and g-SDDMM?
• RQ4 How much porting effort is required to move GNN models between PyTorch, TensorFlow, and MXNet within a single framework-neutral library?

Key findings

• DGL significantly outperforms other popular GNN-oriented frameworks in both speed and memory consumption across benchmarks.
• g-SpMM and g-SDDMM kernels fuse computation and aggregation, reducing memory traffic and enabling training on much larger graphs than PyG can handle (e.g., ML-10m with PyG out of memory).
• On CPU, DGL achieves substantial speedups (1.9x to 64x) over PyG due to higher CPU utilization of its kernels.
• Mini-batch training with cluster sampling (CS) yields about 1.56x speedup for GAT when using DGL.
• DGL provides framework-neutral backends with low overhead, achieving competitive performance relative to framework-specific implementations.