[Paper Review] Fast Graph Representation Learning with PyTorch Geometric

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

Advanced Graph Neural NetworksComputer Science45 references1,259 citations

TL;DR

PyTorch Geometric is a CUDA-accelerated library for fast, flexible graph representation learning built on PyTorch, offering a unified MessagePassing API and extensive operator support.

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.

Motivation & Objective

• Address the challenge of high-throughput GNNs on irregular graph data (graphs, point clouds, manifolds).
• Provide a unified, PyTorch-like API for rapid prototyping of graph neural networks.
• Offer a rich set of neighborhood aggregation, pooling, and dataset processing tools implemented with efficient CUDA kernels.

Proposed method

• Adopt a simple message passing formulation with gather/scatter operations to implement neighborhood aggregation.
• Provide a unified MessagePassing interface where users implement phi (message) and gamma/update (aggregation) with a chosen Dottedsquare (aggregation) function.
• Implement a wide range of GNN operators (GCN, SGC, ARMA, GraphSAGE, GAT, AGNN, GIN, APPNP, DNA, signed operators) and point-cloud/manifold variants (Relational GCN, PointNet++/PointCNN, MPNN, MoNet, SplineCNN, EdgeCNN).
• Support global pooling and hierarchical pooling with multiple readout and pooling options (add/mean/max, Set2Set, SortPooling, global attention, DiffPool, top_k pooling).
• Enable mini-batch processing of variable-sized graphs via sparse block-diagonal adjacency and concatenated feature matrices, with automatic graph assignment to prevent cross-graph leakage.
• Provide easy dataset processing with transforms, automatic dataset downloads, and benchmarks across common graph datasets.

Experimental results

Research questions

• RQ1 How does PyG perform across a broad spectrum of GNN operators on standard benchmarks?
• RQ2 Can gather/scatter based implementations in PyG provide higher throughput on sparse, irregular graphs compared to traditional sparse matrix multiplications?
• RQ3 What is the practicality and reproducibility of PyG’s performance in semi-supervised node classification, graph classification, and point-cloud tasks compared to existing frameworks?
• RQ4 How do various pooling and Jumping Knowledge style readouts affect performance on graph-level tasks?
• RQ5 What runtime advantages does PyG offer relative to other graph learning libraries?

Key findings

• PyG achieves high data throughput by leveraging sparse GPU acceleration and dedicated CUDA kernels.
• The library provides a unified, PyTorch-like API enabling rapid prototyping of new research ideas via the MessagePassing interface.
• A broad set of graph, point-cloud, and manifold operators are implemented, covering many influential GNN architectures.
• Mini-batch processing is supported for graphs of different sizes by constructing a block-diagonal adjacency, enabling scalable training.
• Empirical evaluation shows strong reproducibility and competitiveness across semi-supervised node classification, graph classification, and 3D data processing tasks.
• Compared to DGL, PyG can achieve substantially faster training runtimes on several benchmarks, showcasing practical speed advantages.