[論文レビュー] 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 のメッセージパ Passing を一般化 SpMM（g-SpMM）および一般化 SDDMM（g-SDDMM）演算として形式化する。
• 並列化戦略を開発（g-SpMM はノード並列、g-SDDMM はエッジ並列）し、性能のための CSR/CSC/ COO 形式を検討する。
• DGLGraph を中心データ構造として実装し、フレームワークシムと自動形式切替を用いてフォワードおよびバックワードパスを最適化する。
• フレームワーク変更を最小化してフレームワークに依存しない設計を提供する。テンソル共有には DLPack を使用し、バックプロパゲーションの微分可能な演算子を作成する。
• グラフ中心の API（g.update_all、g.apply_edges）を公開し、ユーザー定義関数と中間テンソルを実体化せずにメッセージパ Passing を組み合わせられるようにする。

実験結果

リサーチクエスチョン

• RQ1 グラフニューラルネットワークの計算は、2つの一般化されたプリミティブ（g-SpMM と g-SDDMM）で効果的に捉えられるのか。
• RQ2 グラフ中心のフレームワーク中立なパッケージは、さまざまな GNN モデルとデータセットにおいて優れた速度とメモリ効率を提供するのか。
• RQ3 g-SpMM および g-SDDMM のために GPU/CPU ハードウェアを最も効果的に活用する並列化戦略は何か。
• RQ4 単一のフレームワーク中立ライブラリ内で、PyTorch、TensorFlow、MXNet 間の GNN モデルの移植作業はどれくらい必要か。

主な発見

• DGL は、ベンチマーク全体で他の人気のある GNN 指向フレームワークよりも速度とメモリ消費の両方で顕著に上回る。
• g-SpMM および g-SDDMM のカーネルは計算と集約を結合してメモリトラフィックを削減し、PyG が扱えるよりはるかに大規模なグラフの学習を可能にする（例：ML-10m を PyG でアウトオブメモリ）。
• CPU では、DGL が PyG と比較してカーネルの CPU 利用率が高いため、1.9x から 64x の大幅な高速化を達成する。
• クラスタサンプリング（CS）を用いたミニバッチ学習は、DGL 使用時の GAT の約 1.56x の高速化をもたらす。
• DGL は低オーバーヘッドのフレームワーク中立バックエンドを提供し、フレームワーク固有の実装と比較して競争力のある性能を達成する。