Getting Started with PyTorch Geometric: Graph Neural Networks in PyTorch
Most machine learning works on grids (images) or sequences (text), but a lot of real-world data is naturally a graph: people connected to friends, atoms bonded into molecules, transactions linking accounts. PyTorch Geometric (PyG) is a library built on top of PyTorch that makes it practical to train neural networks directly on this kind of connected data. This tutorial assumes you are comfortable with basic PyTorch (tensors, nn.Module, a training loop) but have not worked with graph neural networks before.
What Is Graph Machine Learning?
A graph is a set of nodes (also called vertices) connected by edges. Both nodes and edges can carry features. The structure itself — who is connected to whom — carries information that a plain feed-forward network would ignore if you flattened everything into a table.
Graph machine learning shows up in many production settings:
- Social and recommendation graphs. Users, items, and their interactions form a graph. Predicting which item a user will click is a link prediction problem; classifying a user as a bot is node classification.
- Molecules and materials. Atoms are nodes, chemical bonds are edges. Predicting a molecule's solubility or toxicity is a graph-level regression or classification task.
- Fraud and risk. Accounts, devices, and transactions form a graph. Fraud rings are dense subgraphs that are nearly invisible when each transaction is scored independently.
- Knowledge graphs. Entities and typed relations (for example "Paris — capitalof — France") support reasoning and retrieval. These are typically heterogeneous graphs with many node and edge types.
The common theme is that a prediction about one element depends on its neighborhood, not just its own features.
The Intuition: Message Passing
Almost every modern GNN is a form of message passing (sometimes called neighborhood aggregation or graph convolution). The idea is simple and repeats over several layers:
After one layer, every node's representation reflects its immediate neighbors. After two layers, it reflects neighbors-of-neighbors, and so on. Stacking k layers lets information flow across a k-hop neighborhood.
Formally, a generic message passing layer updates node i as:
hi' = UPDATE( hi , AGGREGATE( { MESSAGE(hi, hj, eij) : j in N(i) } ) )
where N(i) are the neighbors of i, and MESSAGE, AGGREGATE, and UPDATE are the parts each GNN variant defines differently. Graph Convolutional Networks (GCN) use a normalized sum, GraphSAGE samples and concatenates, and Graph Attention Networks (GAT) learn weights for each neighbor. We will see these shortly.
Installation
PyG runs on top of PyTorch, so install a matching PyTorch build first. For a CPU-only setup the basics are:
pip install torch
pip install torchgeometric
For most node- and graph-level tutorials this is enough — recent PyG versions implement the core operations in pure PyTorch.
PyG also offers optional compiled extension wheels (pyglib, torchscatter, torchsparse, torchcluster, torchsplineconv). These accelerate scatter/gather and sparse operations and are required for some advanced features (certain samplers and cluster routines). They must match both your PyTorch version and your CUDA build, so install them from the PyG wheel index rather than from PyPI:
# Replace 2.4.0 and cu121 with your own torch version and CUDA tag.
Use cpu instead of cu121 for a CPU-only machine.
pip install pyglib torchscatter torchsparse torchcluster torchsplineconv \
-f https://data.pyg.org/whl/torch-2.4.0+cu121.html
To find the right tags, check your installed versions:
import torch
print(torch.version) # e.g. 2.4.0
print(torch.version.cuda) # e.g. 12.1, or None for CPU
A version mismatch between these wheels and your PyTorch build is the single most common installation problem, so confirm the tags before installing.
The Data Object: One Graph in Memory
A single graph in PyG is a torchgeometric.data.Data object. Its most important attributes are:
x: node feature matrix of shape[numnodes, numnodefeatures].edgeindex: connectivity in COO format, shape[2, numedges]. Columnkmeans "there is an edge fromedgeindex[0, k]toedgeindex[1, k]."edgeattr: optional edge feature matrix, shape[numedges, numedgefeatures].y: targets — one label per node, or a single label for the whole graph.- masks (
trainmask,valmask,testmask): boolean vectors marking which nodes belong to each split.
Let us build a tiny undirected graph with three nodes by hand. Because PyG treats edges as directed, an undirected edge between 0 and 1 is stored as two columns: 0->1 and 1->0.
import torch
from torchgeometric.data import Data
Three nodes, each with a 2-dimensional feature vector.
x = torch.tensor([[1.0, 0.0],
[0.0, 1.0],
[1.0, 1.0]], dtype=torch.float)
Undirected edges 0-1 and 1-2, each listed in both directions.
edgeindex = torch.tensor([[0, 1, 1, 2],
[1, 0, 2, 1]], dtype=torch.long)
One class label per node.
y = torch.tensor([0, 1, 0], dtype=torch.long)
data = Data(x=x, edgeindex=edgeindex, y=y)
print(data)
print("Nodes:", data.numnodes, "Edges:", data.numedges)
print("Has isolated nodes:", data.hasisolatednodes())
The edgeindex tensor must use torch.long, and indices must be valid node ids. PyG provides data.validate() to catch common mistakes early.
Built-in Datasets
PyG ships with many standard benchmark datasets so you can focus on modeling rather than data wrangling.
Cora (viaPlanetoid) is a citation network: nodes are papers, edges are citations, and the task is to classify each paper's topic. It is the canonical node-classification benchmark.
from torchgeometric.datasets import Planetoid
dataset = Planetoid(root="data/Cora", name="Cora")
data = dataset[0] # a single large graph
print("Graphs:", len(dataset))
print("Features per node:", dataset.numnodefeatures)
print("Classes:", dataset.numclasses)
print("Train/val/test nodes:",
int(data.trainmask.sum()),
int(data.valmask.sum()),
int(data.testmask.sum()))
from torchgeometric.datasets import TUDataset
dataset = TUDataset(root="data/MUTAG", name="MUTAG")
print("Number of graphs:", len(dataset))
print("Classes:", dataset.num
classes)
print("First graph:", dataset[0])
Mini-Batching Graphs
For node classification on Cora the whole graph fits in memory, so no batching is needed. For graph classification you have many small graphs, and you batch them with a DataLoader. PyG batches by treating the disjoint union of several graphs as one big graph; a batch vector records which original graph each node came from.
from torchgeometric.loader import DataLoader
loader = DataLoader(dataset, batch
size=32, shuffle=True)
for batch in loader:
print(batch) # a Batch object
print("Graphs in batch:", batch.numgraphs)
print("batch vector shape:", batch.batch.shape) # [totalnodesinbatch]
break
The batch.batch vector is what makes graph-level pooling possible: it tells pooling operators which nodes to aggregate together.
Building a GNN with GCNConv
PyG provides convolution layers as drop-in nn.Modules. The most common ones:
GCNConv— Graph Convolutional Network; a symmetric-normalized neighbor sum.SAGEConv— GraphSAGE; aggregates a (possibly sampled) neighborhood and concatenates with the node's own features, which scales well.GATConv— Graph Attention Network; learns an attention weight per neighbor.GINConv— Graph Isomorphism Network; theoretically very expressive, popular for graph-level tasks.
A two-layer GCN looks almost like an ordinary MLP, except each linear layer is replaced by a graph convolution that also takes edgeindex.
import torch.nn.functional as F
from torch import nn
from torchgeometric.nn import GCNConv
class GCN(nn.Module):
def init(self, inchannels, hiddenchannels, numclasses):
super().init()
self.conv1 = GCNConv(inchannels, hiddenchannels)
self.conv2 = GCNConv(hiddenchannels, numclasses)
def forward(self, x, edgeindex):
x = self.conv1(x, edgeindex)
x = F.relu(x)
x = F.dropout(x, p=0.5, training=self.training)
x = self.conv2(x, edgeindex)
return x # raw logits
Task 1: Node Classification on Cora
We train the GCN to predict each paper's topic, computing the loss only on training nodes and reporting accuracy on each split.
import torch
from torchgeometric.datasets import Planetoid
device = torch.device("cuda" if torch.cuda.isavailable() else "cpu")
dataset = Planetoid(root="data/Cora", name="Cora")
data = dataset[0].to(device)
model = GCN(dataset.numnodefeatures, 16, dataset.numclasses).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weightdecay=5e-4)
def train():
model.train()
optimizer.zerograd()
out = model(data.x, data.edgeindex)
loss = F.crossentropy(out[data.trainmask], data.y[data.trainmask])
loss.backward()
optimizer.step()
return float(loss)
@torch.nograd()
def evaluate():
model.eval()
logits = model(data.x, data.edgeindex)
pred = logits.argmax(dim=1)
accs = {}
for name, mask in [("train", data.trainmask),
("val", data.valmask),
("test", data.testmask)]:
correct = (pred[mask] == data.y[mask]).sum()
accs[name] = float(correct) / int(mask.sum())
return accs
bestval = 0.0
for epoch in range(1, 201):
loss = train()
accs = evaluate()
if accs["val"] > bestval:
bestval = accs["val"]
if epoch % 20 == 0:
print(f"Epoch {epoch:3d} | loss {loss:.4f} | "
f"train {accs['train']:.3f} val {accs['val']:.3f} test {accs['test']:.3f}")
A two-layer GCN typically reaches around 80% test accuracy on Cora. Note the weightdecay on the optimizer — small citation graphs overfit quickly, so regularization matters.
Task 2: Graph Classification on a TUDataset
Here each graph has one label. After message passing produces node embeddings, we collapse them into a single graph embedding with global pooling, then apply a classifier. globalmeanpool uses the batch vector to average node embeddings within each graph.
import torch
from torch import nn
import torch.nn.functional as F
from torchgeometric.datasets import TUDataset
from torchgeometric.loader import DataLoader
from torchgeometric.nn import GCNConv, globalmeanpool
dataset = TUDataset(root="data/MUTAG", name="MUTAG").shuffle()
traindataset = dataset[:150]
testdataset = dataset[150:]
trainloader = DataLoader(traindataset, batchsize=32, shuffle=True)
testloader = DataLoader(testdataset, batchsize=32)
class GraphGCN(nn.Module):
def init(self, inchannels, hidden, numclasses):
super().init()
self.conv1 = GCNConv(inchannels, hidden)
self.conv2 = GCNConv(hidden, hidden)
self.lin = nn.Linear(hidden, numclasses)
def forward(self, x, edgeindex, batch):
x = F.relu(self.conv1(x, edgeindex))
x = F.relu(self.conv2(x, edgeindex))
x = globalmeanpool(x, batch) # [numgraphs, hidden]
x = F.dropout(x, p=0.5, training=self.training)
return self.lin(x)
device = torch.device("cuda" if torch.cuda.isavailable() else "cpu")
model = GraphGCN(dataset.numnodefeatures, 64, dataset.numclasses).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
def trainepoch():
model.train()
total = 0.0
for batch in trainloader:
batch = batch.to(device)
optimizer.zerograd()
out = model(batch.x, batch.edgeindex, batch.batch)
loss = F.crossentropy(out, batch.y)
loss.backward()
optimizer.step()
total += float(loss) batch.numgraphs
return total / len(trainloader.dataset)
@torch.nograd()
def test(loader):
model.eval()
correct = 0
for batch in loader:
batch = batch.to(device)
pred = model(batch.x, batch.edgeindex, batch.batch).argmax(dim=1)
correct += int((pred == batch.y).sum())
return correct / len(loader.dataset)
for epoch in range(1, 101):
loss = trainepoch()
if epoch % 10 == 0:
print(f"Epoch {epoch:3d} | loss {loss:.4f} | "
f"train {test(trainloader):.3f} | test {test(testloader):.3f}")
MUTAG is small, so results vary between runs; averaging over several random splits gives a more honest estimate.
Task 3: Link Prediction
Link prediction asks: given the existing edges, which missing edges are likely real? The standard recipe is encode then decode. A GNN encoder turns nodes into embeddings; a decoder scores a candidate edge (i, j) by the dot product of their embeddings. Because most node pairs are non-edges, we train with negative sampling — pairing each real edge with a randomly drawn non-edge.
import torch
import torch.nn.functional as F
from torchgeometric.datasets import Planetoid
from torchgeometric.nn import GCNConv
from torchgeometric.utils import negativesampling
from torchgeometric.transforms import RandomLinkSplit
device = torch.device("cuda" if torch.cuda.isavailable() else "cpu")
dataset = Planetoid(root="data/Cora", name="Cora")
Split edges into message-passing, supervision, and evaluation sets.
transform = RandomLinkSplit(numval=0.05, numtest=0.1,
isundirected=True, addnegativetrainsamples=False)
traindata, valdata, testdata = transform(dataset[0])
traindata = traindata.to(device)
class Encoder(torch.nn.Module):
def init(self, inchannels, hidden, outchannels):
super().init()
self.conv1 = GCNConv(inchannels, hidden)
self.conv2 = GCNConv(hidden, outchannels)
def forward(self, x, edgeindex):
x = F.relu(self.conv1(x, edgeindex))
return self.conv2(x, edgeindex)
def decode(z, edgelabelindex):
src, dst = edgelabelindex
return (z[src] z[dst]).sum(dim=-1) # dot-product score (logit)
model = Encoder(dataset.numnodefeatures, 128, 64).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
def train():
model.train()
optimizer.zerograd()
z = model(traindata.x, traindata.edgeindex)
# Fresh negative samples each epoch.
negedgeindex = negativesampling(
edgeindex=traindata.edgeindex,
numnodes=traindata.numnodes,
numnegsamples=traindata.edgelabelindex.size(1))
edgelabelindex = torch.cat(
[traindata.edgelabelindex, negedgeindex], dim=1)
edgelabel = torch.cat([
traindata.edgelabel,
traindata.edgelabel.newzeros(negedgeindex.size(1))], dim=0)
out = decode(z, edgelabelindex)
loss = F.binarycrossentropywithlogits(out, edgelabel)
loss.backward()
optimizer.step()
return float(loss)
@torch.nograd()
def test(data):
model.eval()
z = model(traindata.x, traindata.edgeindex)
out = decode(z, data.edgelabelindex).sigmoid()
pred = (out > 0.5).float()
return float((pred == data.edgelabel).float().mean())
for epoch in range(1, 101):
loss = train()
if epoch % 20 == 0:
print(f"Epoch {epoch:3d} | loss {loss:.4f} | "
f"val acc {test(valdata.to(device)):.3f}")
For link prediction, accuracy is a weak metric because of class imbalance; AUC or average precision (from torchmetrics or sklearn) are the usual choices in practice.
Custom Message Passing
When a built-in layer does not fit your needs, subclass MessagePassing. You implement message (how to build a message from a source node to a target) and optionally update, then call self.propagate(...) which handles the gather, the chosen aggregation, and the scatter back to target nodes. The example below re-implements a simple mean-aggregation layer.
import torch
from torch import nn
from torchgeometric.nn import MessagePassing
from torchgeometric.utils import addselfloops
class MeanConv(MessagePassing):
def init(self, inchannels, outchannels):
super().init(aggr="mean") # built-in mean aggregation
self.lin = nn.Linear(inchannels, outchannels)
def forward(self, x, edgeindex):
edgeindex, = addselfloops(edgeindex, numnodes=x.size(0))
x = self.lin(x)
return self.propagate(edgeindex, x=x) # triggers message + aggregate
def message(self, xj):
# xj holds features of source (neighbor) nodes for each edge.
return xj
The j suffix means "source node of the edge"; i would mean the target. aggr can be "mean", "sum", "max", or a custom Aggregation module. This pattern covers most GNN variants you might want to invent.
Heterogeneous Graphs
Many real graphs have multiple node and edge types — for example users and items connected by "buys" and "views" edges. PyG models these with HeteroData, which stores a separate x per node type and a separate edgeindex per relation triple (sourcetype, relation, targettype).
from torchgeometric.data import HeteroData
data = HeteroData()
data["user"].x = torch.randn(100, 16)
data["item"].x = torch.randn(50, 32)
data["user", "buys", "item"].edgeindex = torch.randint(0, 50, (2, 200))
Rather than writing type-specific layers by hand, you can build a normal homogeneous GNN and convert it automatically with tohetero, which replicates the message passing per relation and sums across relations into each target node type.
from torchgeometric.nn import SAGEConv, tohetero
class HGNN(torch.nn.Module):
def init(self, hidden, out):
super().init()
self.conv1 = SAGEConv((-1, -1), hidden) # -1 = infer input dim
self.conv2 = SAGEConv((-1, -1), out)
def forward(self, x, edgeindex):
x = self.conv1(x, edgeindex).relu()
return self.conv2(x, edgeindex)
model = HGNN(hidden=64, out=8)
model = tohetero(model, data.metadata(), aggr="sum")
out = model(data.xdict, data.edgeindexdict)
Using (-1, -1) and lazy initialization lets the model infer per-type input dimensions on the first forward pass.
Scaling to Large Graphs with Neighbor Sampling
Full-batch training (the Cora example) loads the entire graph each step. That breaks down on graphs with millions of nodes. NeighborLoader instead samples a fixed number of neighbors per hop around a batch of target nodes, producing subgraphs small enough to fit in GPU memory. This is the GraphSAGE training strategy.
from torchgeometric.loader import NeighborLoader
loader = NeighborLoader(
data,
numneighbors=[15, 10], # 15 neighbors in hop 1, 10 in hop 2
batchsize=128,
inputnodes=data.trainmask,
shuffle=True,
)
for subgraph in loader:
# The first batchsize nodes are the seed (target) nodes.
print(subgraph)
break
You then train on each sampled subgraph as if it were a small full graph, computing the loss on the seed nodes only. Pairing NeighborLoader with SAGEConv is a well-tested combination for large-scale node classification.
Practical Tips
Over-smoothing. Stacking too many message-passing layers makes every node's embedding converge to the same value, so accuracy drops. Most graphs do best with two to four layers. If you need depth, use residual connections, jumping-knowledge concatenation, orGCNConv with addselfloops.
Normalization. Add BatchNorm or LayerNorm between layers (torchgeometric.nn.BatchNorm is graph-aware). It stabilizes training and lets you go slightly deeper.
Self-loops and direction. GCNConv adds self-loops by default so a node keeps its own information; check whether other layers do. For directed graphs, decide deliberately whether to add reverse edges — many tasks treat the graph as undirected.
Evaluation. Use the dataset's official splits when they exist. For link prediction and imbalanced node tasks, report AUC or average precision rather than accuracy. Average over several seeds, since small graphs have high variance.
Reproducibility. Set seeds with torchgeometric.seedeverything(seed) and remember that sampling-based loaders introduce randomness across runs.
Sanity checks. Call data.validate(), confirm edgeindex is long and within bounds, and check for isolated nodes before blaming the model for poor results.
Conclusion and Key Takeaways
- Graph machine learning applies when relationships between data points matter — social networks, molecules, fraud rings, knowledge graphs.
- Almost every GNN is message passing: aggregate neighbor features, combine with a node's own features, repeat across layers so information flows over multi-hop neighborhoods.
- The
Dataobject holds one graph (x,edgeindexin COO,edgeattr,y, masks).HeteroDataextends this to multiple node and edge types. - PyG's built-in datasets (
Planetoid,TUDataset) andDataLoadermini-batching (thebatchvector) let you start modeling immediately. - The three core tasks — node classification, graph classification (with
globalmean_pool), and link prediction (encode/decode with negative sampling) — share the same convolution layers likeGCNConv,SAGEConv,GATConv, andGINConv. - For anything custom, subclass
MessagePassing; for scale, useNeighborLoader; and watch for over-smoothing by keeping networks shallow and adding normalization.
With these building blocks you can move from the Cora toy example to production graphs while reusing the same patterns.