[論文レビュー] Fast Graph Representation Learning with PyTorch Geometric

Matthias Fey, Jan Eric Lenssen|arXiv (Cornell University)|Mar 6, 2019

Advanced Graph Neural Networks参考文献 45被引用数 1,259

ひとこと要約

PyTorch Geometric は、PyTorch 上に構築された CUDA 加速の高速かつ柔軟なグラフ表現学習ライブラリであり、統一された MessagePassing API と豊富な演算子サポートを提供します。

ABSTRACT

We introduce PyTorch Geometric, a library for deep learning on irregularly structured input data such as graphs, point clouds and manifolds, built upon PyTorch. In addition to general graph data structures and processing methods, it contains a variety of recently published methods from the domains of relational learning and 3D data processing. PyTorch Geometric achieves high data throughput by leveraging sparse GPU acceleration, by providing dedicated CUDA kernels and by introducing efficient mini-batch handling for input examples of different size. In this work, we present the library in detail and perform a comprehensive comparative study of the implemented methods in homogeneous evaluation scenarios.

研究の動機と目的

• 不規則なグラフデータ（グラフ、点群、多様体）に対する高スループットGNNの課題に対処する。
• グラフニューラルネットワークの迅速なプロトタイピングのための、統一されたPyTorchライクなAPIを提供する。
• 効率的なCUDAカーネルで実装された豊富な近傍集約、プーリング、データセット処理ツールを提供する。

提案手法

• 集合集散操作を用いたシンプルなメッセージパッシングの定式化を採用し、近傍集約を実装する。
• phi (message) と gamma/update (aggregation) を選択された Dottedsquare (aggregation) 関数とともに実装する、統一された MessagePassing インターフェースを提供する。
• 多数のGNN演算子（GCN, SGC, ARMA, GraphSAGE, GAT, AGNN, GIN, APPNP, DNA, signed operators）および点群/多様体のバリアント（Relational GCN, PointNet++/PointCNN, MPNN, MoNet, SplineCNN, EdgeCNN）を実装する。
• 複数のリードアウトとプーリングオプションを備えたグローバルプーリングおよび階層プーリングをサポートする（add/mean/max, Set2Set, SortPooling, global attention, DiffPool, top_k pooling）。
• 疎ブロック対角隣接行列と連結された特徴行列を用いた可変サイズグラフのミニバッチ処理を有効にし、グラフ間のリークを防ぐ自動的なグラフ割り当てを行う。
• 変換による容易なデータセット処理、データセットの自動ダウンロード、一般的なグラフデータセットにおけるベンチマークを提供する。

Figure 1: Computation scheme of a GNN layer by leveraging gather and scatter methods based on edge indices $\bm{I}$ , hence alternating between node parallel space and edge parallel space.

実験結果

リサーチクエスチョン

• RQ1 標準ベンチマークで、広範なGNN演算子に対してPyGはどう性能を示すか。
• RQ2 PyGの gather/scatter ベースの実装は、従来の疎行列乗算と比較して、スパースで不規則なグラフ上でより高いスループットを提供できるか。
• RQ3 半教師ありノード分類、グラフ分類、点群タスクにおけるPyGの性能の実用性と再現性は、既存フレームワークと比較してどうか。
• RQ4 さまざまなプーリングとJumping Knowledgeスタイルのリードアウトが、グラフレベルのタスクの性能にどう影響しますか。
• RQ5 他のグラフ学習ライブラリと比較して、PyGが提供する実行時の利点は何ですか。

主な発見

• PyG は疎GPUアクセラレーションと専用CUDAカーネルを活用して高いデータスループットを実現します。
• このライブラリは統一されたPyTorchライクなAPIを提供し、MessagePassingインターフェースを介して新しい研究アイデアの迅速なプロトタイピングを可能にします。
• 広範なグラフ、点群、多様体の演算子が実装されており、多くの影響力のあるGNNアーキテクチャを網羅します。
• 異なるサイズのグラフのミニバッチ処理は、ブロック対角隣接行列を構築することで実現され、スケーラブルな訓練を可能にします。
• 実証評価は、半教師ありノード分類、グラフ分類、3Dデータ処理タスクにおいて強い再現性と競争力を示します。
• DGLと比較して、複数のベンチマークで学習実行時間を大幅に短縮できることが示され、実用的な速度利点を示します。

Figure 2: Forward and backward runtimes of $1\,000$ runs of gather and scatter operations (GS) in comparison to sparse-matrix multiplications (SpMM) on Erdős Rényi graphs with $10\,000$ nodes and varying average node degrees. Runtimes are capped at two seconds due to visualization. We report runtime