PyG Practitioner's Guide

PyTorch Geometric (PyG) is a library for graph machine learning built on PyTorch. Created by Matthias Fey, founding engineer at Kumo.ai, it provides graph data structures, message-passing layers, datasets, and tooling for graph learning in research and production.

Get started in five minutes

PyG requires PyTorch. Install both packages, then verify that PyG imports correctly.

pip install torch torch-geometric
# Verify installation
python -c "import torch_geometric; print(torch_geometric.__version__)"

The Cora citation network is a useful first node-classification example. It contains 2,708 papers, 10,556 citations, and seven research-topic classes.

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
dataset = Planetoid(root="/tmp/Cora", name="Cora")
data = dataset[0]
class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)
    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.5, training=self.training)
        return self.conv2(x, edge_index)
model = GCN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(200):
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

Evaluate the trained model on the test split.

model.eval()
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1)
correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
accuracy = correct / data.test_mask.sum()
print(f"Test accuracy: {accuracy:.4f}")

The source guide reports approximately 81% test accuracy for this basic two-layer GCN on Cora.

From relational tables to predictions

Every relational database is a graph hiding in plain sight. Rows are nodes, and foreign-key relationships are edges. PyG makes that structure explicit so a model can learn from both record attributes and connections.

Choose a GNN layer

PyG includes a broad collection of layers because different graph structures and operational constraints require different message-passing strategies.

Start here

Attention, heterogeneous data, and transformers

For the foundations behind message passing, graph types, and graph transformers, see Introduction to Graph Transformers.

Concepts to understand

The following ideas determine whether a graph model is well matched to the problem.

• Message passing: The basic computation through which nodes exchange and aggregate neighbor information.
• Heterogeneous graphs: Graphs with multiple node and edge types, typical of relational databases.
• Over-smoothing: A failure mode where many GNN layers make node representations too similar.
• Link prediction: Predicting missing or future edges for recommendations, fraud detection, and knowledge graphs.
• Graph transformers: Architectures that combine graph structure with broader attention patterns.
• Neighbor sampling: Training on a bounded subset of each node’s neighborhood to scale to large graphs.
• Data leakage: Accidentally using future or otherwise unavailable graph information at training time.

Apply graph learning to business problems

Graph ML is useful when relationships add important signal beyond a single record.

Take graph models to production

Moving from a notebook to a production system requires more than a trained model. Teams need a reliable way to construct and update the graph, train at the needed scale, score new data, handle temporal relationships correctly, and explain predictions.

Common production concerns include:

• Heterogeneous graphs: Model multiple table types and typed relationships.
• Temporal graphs: Avoid using information that would not have been available at prediction time.
• Scaling: Use sampling and partitioning for graphs with millions or billions of nodes.
• Serving: Turn trained embeddings and model outputs into reliable predictions.
• Integration: Connect graph learning workflows to data platforms such as Snowflake and Databricks.
• Explainability: Inspect which graph features and relationships contributed to a prediction.

Datasets for benchmarking

PyG provides loaders for standard graph-learning datasets across citation, social, molecular, recommendation, knowledge-graph, financial, and large-scale benchmark domains.

Examples include Cora, CiteSeer, PubMed, Reddit, OGB-Products, OGB-Papers100M, MUTAG, QM9, MovieLens 1M, Elliptic Bitcoin, and FB15k-237. Choose a dataset that matches the learning task and graph properties you need to evaluate, such as node classification, link prediction, graph classification, or regression.

Further reading

• PyTorch Geometric documentation
• PyTorch Geometric on GitHub
• Graph Convolutional Networks
• GraphSAGE
• Graph Attention Networks
• Graph Transformers