Source-linked AI summary
Fast Graph Representation Learning with PyTorch Geometric
Matthias Fey, Jan Eric Lenssen
TL;DR
Implementing graph neural networks efficiently is challenging because graph data are sparse, irregular, and variable-sized. PyTorch Geometric provides a unified PyTorch-based framework for such data, and experiments show high reproducibility alongside fast training, with models running up to 40 times faster than a DGL baseline.
Problem
Efficiently implementing graph neural networks is challenging because input data are sparse, irregular, and variable-sized.
Method
PyTorch Geometric provides a unified PyTorch-based framework with graph representations, message-passing operators, and processing methods for graphs, point clouds, and manifolds.
Results
Experiments show high reproducibility across most methods, while PyG trains models up to 40 times faster than DGL's Degree Bucketing approach.
Takeaways & Limitations
PyTorch Geometric offers a fast framework for representation learning on graphs, point clouds, and manifolds.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Graph Neural Networks learn localized representations on graphs, point clouds, and manifolds by passing, transforming, and aggregating information. PyTorch Geometric addresses implementation challenges from sparse, irregular, variable-sized data through dedicated CUDA kernels and a unified message passing framework.
- GNNs extract localized embeddings on graphs, point clouds, and manifolds through hierarchical information passing, transformation, and aggregation.
- PyG is a geometric deep learning extension library for PyTorch designed to address high-throughput challenges on sparse, irregular data of varying size.
- Dedicated CUDA kernels enable PyG to achieve high performance, while a simple message passing API unifies recently proposed convolutional and pooling layers.
2 OVERVIEW
PyG provides PyTorch-inspired graph representations and APIs for message passing, pooling, batching, and dataset processing. Its overview emphasizes extensible operators and efficient GPU-compatible handling of irregular graph data.
- Graph representation: PyG represents graphs with node features and sparse COO edge indices, optionally storing multidimensional edge features, while keeping user-facing APIs familiar through PyTorch-inspired design.The representation uses X ∈ ℝ^(N×F), I ∈ ℕ^(2×E), and optional E ∈ ℝ^(E×D).
- Neighborhood aggregation: Its MessagePassing interface lets users define message and update functions plus an aggregation scheme, with source and target node features mapped automatically.Neighborhood aggregation is implemented through gathering and scattering node features with vectorized element-wise computation, enabling GPU acceleration.
- Implemented operators: PyG integrates neighborhood aggregation operators for arbitrary graphs, point clouds, manifolds, and graphs with multidimensional edge features, including GCN, GraphSAGE, GAT, PointNet++, SplineCNN, and EdgeCNN.The overview also lists spectral filters, relational GCN, PointCNN, MPNN, MoNet, and other operators.
- High-level methods: Beyond operators, PyG provides high-level implementations for mutual-information maximization, graph autoencoding, jumping-knowledge aggregation, and temporal-event prediction in knowledge graphs.These implementations extend the library beyond basic neighborhood aggregation layers.
- Pooling: PyG supports graph-level outputs through global add, mean, max, set-to-set, sort, and global soft-attention pooling, alongside hierarchical pooling methods for deeper GNN models.Hierarchical examples include Graclus, voxel-grid pooling, and farthest-point sampling followed by k-NN or query-ball graph generation.
- Batching and datasets: For batches of differently sized graphs, PyG constructs a sparse block-diagonal adjacency matrix, concatenates node features, and generates assignment vectors without modifying neighborhood aggregation methods.The framework also provides a consistent dataset format, transforms, automatic downloading and processing, and over 60 graph-kernel benchmark datasets.
3 EMPIRICAL EVALUATION
The evaluation tests PyTorch Geometric methods across node classification, graph classification, point-cloud classification, and training-runtime experiments under standardized, reproducible setups. Results are generally reproducible, with APPNP strongest in node classification, near-equivalent point-cloud methods, and PyG up to 40 times faster than DGL’s Degree Bucketing approach.
- Evaluation setup: Experiments closely follow the respective papers’ hyperparameters, use homogeneous evaluation scenarios, and are reproducible from the provided GitHub code.Dataset descriptions and statistics are provided in Appendix B.
- Semi-supervised Node Classification: Node-classification results are averaged over 100 runs using both a fixed split and randomly initialized splits with uniform training-class distribution.Test performance is worse for all models with random data splits, while APPNP generally performs best; ARMA, SGC, GCN, and GAT follow closely.
- Graph Classification: Graph classification uses 10-fold cross-validation, discrete node features, degree one-hot encodings when needed, and global mean graph-level outputs.Under standardized evaluations, pooling operators generally underperform their flat counterparts except DiffPool, especially against Jumping Knowledge-enhanced baselines.
- Point Cloud Classification: Point-cloud methods on ModelNet10 perform nearly identically when using approximately equal parameter counts, with PointCNN taking a slight lead.The authors attribute the similarity to shared operating principles and potentially equal expressive power on the task.
- Runtime Experiments: 40 times faster: PyG trains models by up to this factor compared with DGL v0.2’s Degree Bucketing approach in 200-epoch runs on a single NVIDIA GTX 1080 Ti.The experiments measure the runtime of complete training procedures across multiple dataset-model pairs and report that PyG is fast despite sparse data.
4 ROADMAP AND CONCLUSION
The paper presents PyTorch Geometric as a fast representation-learning framework for graphs, point clouds, and manifolds, with plans to expand its integrated method collection through collaboration.
- Roadmap and Conclusion: PyTorch Geometric supports fast representation learning on graphs, point clouds, and manifolds.The framework is designed to integrate existing methods and rapidly incorporate future methods.
- Roadmap and Conclusion: The developers invite researchers and software engineers to collaborate in extending the framework’s scope.This collaboration is part of the project’s ongoing expansion roadmap.
A GATHER AND SCATTER OPERATIONS
PyTorch Geometric uses gather and scatter operations to map node and edge information efficiently between parallel spaces, with custom mean and max operations. Compared with sparse-matrix multiplication, these operations remain fast across input ordering but can throttle on high-degree graphs and are GPU-nondeterministic.
- Implementation: Gather and scatter operations map node and edge information into edge- and node-parallel spaces, using parallelization and atomic operations for high throughput.The implementation provides custom mean and max operations following PyTorch’s scatter_add.
- Runtime comparison: Atomic operations begin throttling runtime only for graphs with approximately 128 average node degree, while exceeding highly optimized SpMM for forward and backward passes.Figure 2 compares gather/scatter runtime with sparse-matrix multiplication.
- Runtime comparison: Unlike SpMM, gather and scatter remain fast regardless of input coalescing and avoid the GPU expense of computing coalesced sparse tensors during preprocessing.SpMM converts adjacency matrices to CSR format and coalesces during backward propagation when transposing the sparse tensor.
- Limitations: Gather and scatter could be improved for highly dense graphs and lower memory use, while GPU scatter operations remain inherently nondeterministic.Inference showed no observed deviations, but training results can vary across identical manual seeds.
B DATASETS
The experiments use citation-network, social-network, bioinformatic, and 3D-object datasets, with detailed descriptions and statistics provided for these datasets. The tasks span document classification, graph-level social and molecular/protein classification, and object-category classification.
- Citation Networks: Cora, Citeseer, and Pubmed represent documents as nodes with bag-of-words features and citation links as undirected edges.Training uses 20 labels per class.
- Social Network Datasets: COLLAB labels researcher ego-networks by scientific field, while IMDB-BINARY uses actor ego-networks to infer movie genre.In IMDB-BINARY, actors or actresses are connected when they appear in the same movie.
- Bioinformatic Datasets: MUTAG contains mutagenetic aromatic and heteroaromatic nitro compounds, while PROTEINS represents proteins as graphs of connected secondary structure elements.PROTEINS connects secondary structure elements when they are neighbors in the amino acid sequence or 3D space.
- 3D Object Datasets: ModelNet10 contains orientation-aligned CAD models, each assigned to exactly one of 10 object categories.The categories were selected from common object categories in the world.