README.md

Examples

This folder contains a plethora of examples covering different GNN use-cases. This readme highlights some key examples.

Note

We recommend the NVIDIA PyG Container for best results and easiest setup with NVIDIA GPUs. See the cuGraph installation guide for details.

A great and simple example to start with is gcn.py, showing a user how to train a GCN model for node-level prediction on small-scale homogeneous data.

For a simple GNN based link prediction example, see link_pred.py.

For an improved GNN based link prediction approach using Attract-Repel embeddings that can significantly boost accuracy (up to 23% improvement in AUC), see ar_link_pred.py. This approach is based on Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs.

To see an example for doing link prediction with an advanced Graph Transformer called LPFormer, see [lpformer.py].

For examples on Open Graph Benchmark datasets, see the ogbn_*.py examples:

• ogbn_train.py is an example for training a GNN on the large-scale ogbn-papers100m dataset, containing approximately ~1.6B edges or the medium scale ogbn-products dataset, ~62M edges.
  • Uses SGFormer (a kind of GraphTransformer) by default.
  • SGFormer Paper
  • Polynormer
  • EGT
  • Kumo.ai x NVIDIA x Stanford Graph Transformer Webinar
• ogbn_proteins_deepgcn.py is an example to showcase how to train deep GNNs on the ogbn-proteins dataset.
• ogbn_train_perforatedai.py shows how to optimize the ogbn_train.py workflow using Perforated AI. Perforated AI provides a PyTorch add-on which increases network accuracy by empowering each artificial neuron with artificial dendrites.

For an example on Relational Deep Learning with the RelBench datasets, see rdl.py.

For examples on using torch.compile, see the examples under examples/compile.

For examples on scaling PyG up via multi-GPUs, see the examples under examples/multi_gpu.

For examples on working with heterogeneous data, see the examples under examples/hetero.

For examples on co-training LLMs with GNNs, see the examples under examples/llm.

• Stanford GNN+LLM Talk

We recommend looking into PyTorch documentation for examples on setting up model parralel GNNs.

Scale to Trillions of Edges with cuGraph

cuGraph is a collection of packages focused on GPU-accelerated graph analytics including support for property graphs and scaling up to thousands of GPUs. cuGraph supports the creation and manipulation of graphs followed by the execution of scalable fast graph algorithms. It is part of the RAPIDS accelerated data science framework.

cuGraph GNN is a collection of GPU-accelerated plugins that support PyTorch and PyG natively through the cuGraph-PyG and WholeGraph subprojects. cuGraph GNN is built on top of cuGraph, leveraging its low-level pylibcugraph API and C++ primitives for sampling and other GNN operations (libcugraph). It also includes the libwholegraph and pylibwholegraph libraries for high-performance distributed edgelist and embedding storage. Users have the option of working with these lower-level libraries directly, or through the higher-level API in cuGraph-PyG that directly implements the GraphStore, FeatureStore, NodeLoader, and LinkLoader interfaces.

Complete documentation on RAPIDS graph packages, including cugraph, cugraph-pyg, pylibwholegraph, and pylibcugraph is available on the RAPIDS docs pages.

See rapidsai/cugraph-gnn/tree/branch-25.12/python/cugraph-pyg/cugraph_pyg/examples on GitHub for fully scalable PyG example workflows.