PyTorch Geometric (PyG)

PyG is the standard library for Graph Neural Networks built on PyTorch. It provides data structures for graphs, 60+ GNN layer implementations, scalable mini-batch training, and support for heterogeneous graphs.

Installation

Reviewed released torch-geometric 2.8.0.post1 (2026-10-01); CPU examples tested with Python 3.13 / PyTorch 2.14.1. Rolling latest docs identify 2.9.0; check the installed release before adopting new APIs. PyG 2.8 requires PyTorch 2.9+; its original release table covers 2.9–2.12. Our 2.14.1 core tests do not establish every extension/backend combination.

# Install the PyTorch build for your platform from https://pytorch.org/get-started/locally/
uv pip install torch==2.14.1
uv pip install torch-geometric==2.8.0.post1
python -c "import torch, torch_geometric; print(torch.__version__, torch.version.cuda, torch_geometric.__version__)"

Basic tensor-based layers need no extensions. Neighbor sampling requires pyg-lib or torch-sparse; spatial k-NN operators require pyg-lib in 2.8. torch-cluster and torch-spline-conv are deprecated and ignored. Inspect the wheel index for your exact Python/OS/Torch/CUDA tuple. Never install wheels for a different Torch release merely because core imports succeed. The tested macOS ARM CPU extension is below; choose a different matching wheel for other platforms, and verify its operators:

uv pip install --only-binary=:all: "pyg-lib==0.9.0+pt214" \
  -f https://data.pyg.org/whl/torch-2.14.0+cpu.html

Conda packages are no longer provided for Torch >2.5. See installation and 2.8 release changes. Optional sampling/GPU/distributed/download examples below are illustrative unless covered by the CPU checks in review notes.

Core Concepts

Graph Data: Data and HeteroData

A graph lives in a Data object. The key attributes:

from torch_geometric.data import Data

data = Data(
    x=node_features,          # [num_nodes, num_node_features]
    edge_index=edge_index,     # [2, num_edges] — COO format, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]
    y=labels,                  # node-level [num_nodes, *] or graph-level [1, *]
    pos=positions,             # [num_nodes, num_dimensions] (for point clouds/spatial)
)

**edge_index format is critical**: it's a [2, num_edges] tensor where edge_index[0] = source nodes, edge_index[1] = target nodes. It is NOT a list of tuples. If you have edge pairs as rows, transpose and call .contiguous():

# If edges are [[src1, dst1], [src2, dst2], ...] — transpose first:
edge_index = edge_pairs.t().contiguous()

For undirected graphs, include both directions: edge (0,1) needs both [0,1] and [1,0] in edge_index.

If node features are absent, set data.num_nodes explicitly from the node table. Inferring it from edge_index.max() + 1 misses isolated nodes, which can corrupt batching offsets and outputs. Check data.validate(raise_on_error=True) after construction, including an edge-free or isolated-node case.

For heterogeneous graphs, use HeteroData — see the Heterogeneous Graphs section below.

Datasets

PyG bundles many standard datasets that auto-download and preprocess:

from torch_geometric.datasets import Planetoid, TUDataset

# Single-graph node classification (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data/Cora', name='Cora', split='public')
data = dataset[0]  # single graph with train/val/test masks

# Multi-graph classification (ENZYMES, MUTAG, IMDB-BINARY, etc.)
dataset = TUDataset(root='./data/TU', name='ENZYMES')
# dataset[0], dataset[1], ... are individual graphs

Common datasets by task:

• Node classification: Planetoid (Cora/Citeseer/Pubmed), OGB (ogbn-arxiv, ogbn-products, ogbn-mag)
• Graph classification: TUDataset (MUTAG, ENZYMES, PROTEINS, IMDB-BINARY), OGB (ogbg-molhiv)
• Link prediction: OGB (ogbl-collab, ogbl-citation2)
• Molecular: QM7b, QM9, MoleculeNet
• Point cloud/mesh: ShapeNet, ModelNet(name="10" or "40"), FAUST (manual download)

Dataset classes manage provider downloads; they are not API search endpoints. Preserve the dataset version, split and preprocessing. OGB benchmarks use the separate ogb package/evaluator; do not replace their official split with a random split. See review notes for verified download locations and unexecuted large datasets.

Transforms

Transforms preprocess or augment graph data, analogous to torchvision transforms:

import torch_geometric.transforms as T
from torch_geometric.datasets import ShapeNet

# Common transforms
T.NormalizeFeatures()    # Shift by minimum, then divide row sum (clamped >=1)
T.ToUndirected()         # Add reverse edges to make graph undirected
T.AddSelfLoops()         # Add self-loop edges
T.KNNGraph(k=6)          # Build k-NN graph from positions; requires pyg-lib
T.RandomJitter(0.01)     # Random noise augmentation on positions
T.Compose([...])         # Chain multiple transforms

# Apply as pre_transform (once, saved to disk) or transform (every access)
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

ToUndirected may merge/reduce duplicate edge attributes (default sum); confirm weight/label semantics before applying it. Adding self-loops can also duplicate existing loops. Do not make directed or temporal relations undirected without a scientific reason.

Building GNN Models

Quick Start: Using Built-in Layers

The fastest way to build a GNN — stack conv layers from torch_geometric.nn:

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

GCNConv, SAGEConv and attention layers return embeddings; add the intended nonlinearities between them. GINConv/EdgeConv use supplied networks that may already contain activations.

Choosing a Conv Layer

Pick based on your task and graph structure:

Check the chosen signature: RGCNConv also needs relation IDs (edge_type), HGTConv takes dictionaries, and GCNConv accepts scalar edge_weight, not arbitrary edge_attr.

Lazy Initialization

Use -1 for input channels to let PyG infer dimensions automatically — especially useful for heterogeneous models:

from torch_geometric.nn import SAGEConv
conv = SAGEConv((-1, -1), 64)  # Input dims inferred on first forward pass
# Initialize lazy modules:
with torch.no_grad():
    out = conv(data.x, data.edge_index)

High-Level Model APIs

For common architectures, PyG provides ready-made model classes:

from torch_geometric.nn import GraphSAGE, GCN as GCNModel, GAT as GATModel, GIN as GINModel

model = GraphSAGE(
    in_channels=dataset.num_features,
    hidden_channels=64,
    out_channels=dataset.num_classes,
    num_layers=2,
)

Custom Layers via MessagePassing

To implement a novel GNN layer, subclass MessagePassing. The framework is:

1. propagate() orchestrates the message passing
2. message() defines what info flows along each edge (the phi function)
3. aggregate() combines messages at each node (sum/mean/max)
4. update() transforms the aggregated result (the gamma function)

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class MyConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # "add", "mean", or "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # Pre-processing before message passing
        x = self.lin(x)
        # Start message passing
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: features of source nodes for each edge [num_edges, features]
        # The _j suffix auto-indexes source nodes, _i indexes target nodes
        return x_j

**The _i / _j convention**: any tensor passed to propagate() can be auto-indexed by appending _i (target/central node) or _j (source/neighbor node) in the message() signature. So if you pass x=... to propagate, you can access x_i and x_j in message().

Read references/message_passing.md for the full GCN and EdgeConv implementation examples.

Task-Specific Patterns

Training loops are adaptation recipes. Regression checks use tiny synthetic inputs and short runs, not full benchmark convergence.

Node Classification

# Full-batch training on a single graph (e.g., Cora)
model = GCN(dataset.num_features, 64, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
model.train()
for epoch in range(200):
    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()

# Select checkpoints using validation only; evaluate test once afterward.
model.eval()  # Module evaluation behavior; gradients are disabled separately.
with torch.no_grad():
    pred = model(data.x, data.edge_index).argmax(dim=1)
    acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()

Graph Classification

Multiple graphs — use DataLoader for mini-batching and global pooling to get graph-level representations:

from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool

loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

class GraphClassifier(torch.nn.Module):
    def __init__(self, in_ch, hidden_ch, out_ch):
        super().__init__()
        self.conv1 = GCNConv(in_ch, hidden_ch)
        self.conv2 = GCNConv(hidden_ch, hidden_ch)
        self.lin = torch.nn.Linear(hidden_ch, out_ch)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        x = global_mean_pool(x, batch)  # [num_graphs_in_batch, hidden_ch]
        return self.lin(x)

# train_dataset is a previously split graph-level dataset with node features.
model = GraphClassifier(dataset.num_features, 64, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
model.train()
for data in loader:
    optimizer.zero_grad()
    out = model(data.x, data.edge_index, data.batch)
    loss = F.cross_entropy(out, data.y.view(-1).long())
    loss.backward()
    optimizer.step()

PyG's DataLoader offsets edge indices to represent a disconnected union (block-diagonal adjacency), without allocating a dense matrix. The batch tensor maps each node to its graph index. Pooling ops (global_mean_pool, global_max_pool, global_add_pool) use this to aggregate per-graph.

Link Prediction

Split edges into train/val/test, use negative sampling:

from torch_geometric.transforms import RandomLinkSplit

transform = RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    is_undirected=True,
    add_negative_train_samples=True,
    disjoint_train_ratio=0.2,  # Keep supervision out of training message edges.
)
train_data, val_data, test_data = transform(data)

# Encode nodes, then score edges
encoder = GCN(data.num_features, 64, 32)
z = encoder(train_data.x, train_data.edge_index)
src, dst = train_data.edge_label_index
logits = (z[src] * z[dst]).sum(dim=-1)
loss = F.binary_cross_entropy_with_logits(logits, train_data.edge_label.float())

Read references/link_prediction.md for the complete link prediction guide: GAE/VGAE autoencoders, full training loops, LinkNeighborLoader for large graphs, heterogeneous link prediction, and evaluation metrics.

Scaling to Large Graphs

For graphs that don't fit in GPU memory, use NeighborLoader with a compatible sampling backend. The following is illustrative; it needs pyg-lib or torch-sparse:

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[15, 10],     # Sample 15 neighbors in hop 1, 10 in hop 2
    batch_size=128,              # Number of seed nodes per batch
    input_nodes=data.train_mask, # Which nodes to sample from
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # Only use first batch_size nodes for loss (these are the seed nodes)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

Key points about NeighborLoader:

• num_neighbors list length should match GNN depth (number of message passing layers)
• Seed nodes are always the first batch.batch_size nodes in the output
• batch.n_id maps relabeled indices back to original node IDs
• Works for both Data and HeteroData
• For link prediction, use LinkNeighborLoader instead
• Large fan-out across many hops grows rapidly; measure the sampled sizes

Other scalability options: ClusterLoader (ClusterGCN), GraphSAINTSampler, ShaDowKHopSampler. For multi-GPU training, DDP, PyTorch Lightning integration, and torch.compile support, read references/scaling.md.

Heterogeneous Graphs

For graphs with multiple node and edge types (social networks, knowledge graphs, recommendation):

from torch_geometric.data import HeteroData

data = HeteroData()

# Node features — indexed by node type string
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)

# Edge indices — indexed by (src_type, edge_type, dst_type) triplet
data['user', 'rates', 'movie'].edge_index = torch.stack([
    torch.randint(1000, (3000,)), torch.randint(500, (3000,))])
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))

# Access convenience dicts
data.x_dict        # {'user': tensor, 'movie': tensor}
data.edge_index_dict  # {('user','rates','movie'): tensor, ...}
data.metadata()    # ([node_types], [edge_types])

Three ways to build heterogeneous GNNs

**1. Auto-convert with to_hetero()** — write a homogeneous model, convert automatically:

from torch_geometric.nn import SAGEConv, to_hetero

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

model = GNN(64, dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

# Now accepts dicts:
out = model(data.x_dict, data.edge_index_dict)

Use (-1, -1) for bipartite input channels (source, target may differ). Lazy init handles the rest.

**2. HeteroConv wrapper** — different conv per edge type:

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv

conv = HeteroConv({
    ('paper', 'cites', 'paper'): GCNConv(-1, 64),
    ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
    ('paper', 'rev_writes', 'author'): GATConv((-1, -1), 64, add_self_loops=False),
}, aggr='sum')

3. Native heterogeneous operators like HGTConv:

from torch_geometric.nn import HGTConv
conv = HGTConv(-1, 64, data.metadata(), heads=4)  # 64 must divide by 4

Important for heterogeneous graphs:

• Use T.ToUndirected() to add reverse edge types for bidirectional message flow
• Disable add_self_loops in bipartite conv layers (different source/dest types) — use skip connections instead: conv(x, edge_index) + lin(x)
• For NeighborLoader on HeteroData, specify input_nodes as ('node_type', mask) tuple
• num_neighbors can be a dict keyed by edge type for fine-grained control

Read references/heterogeneous.md for complete examples including training loops and NeighborLoader usage with heterogeneous graphs.

Custom Datasets

For loading your own data into PyG:

• Quick (no class needed): Create Data objects directly and pass a list to DataLoader
• Reusable (fits in RAM): Subclass InMemoryDataset — override raw_file_names, processed_file_names, download(), process()
• Large (disk-backed): Subclass Dataset — also override len() and get()
• From CSV: Load node/edge tables with pandas, build mappings to consecutive indices, assemble into Data or HeteroData
• From NetworkX: from_networkx(G) converts a NetworkX graph directly
• From scipy sparse: from_scipy_sparse_matrix(adj) extracts edge_index

Read references/custom_datasets.md for complete examples with all patterns, CSV loading with encoders, and the MovieLens walkthrough.

Explainability

PyG provides torch_geometric.explain for interpreting GNN predictions:

from torch_geometric.explain import Explainer, GNNExplainer

explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config=dict(
        mode='multiclass_classification',
        task_level='node',
        return_type='raw',  # GCN above returns logits.
    ),
)

explanation = explainer(data.x, data.edge_index, index=10)
explanation.visualize_graph()           # Important subgraph
explanation.visualize_feature_importance(top_k=10)  # Feature importance

Available algorithms: GNNExplainer (optimization-based), PGExplainer (parametric, trained), CaptumExplainer (gradient-based via Captum), AttentionExplainer (attention weights). Heterogeneous support depends on the algorithm; wrap dict-returning models to select one output node type.

Read references/explainability.md for all algorithms, heterogeneous explanations, evaluation metrics, and PGExplainer training.

Common Pitfalls

1. edge_index shape: Must be [2, num_edges], not [num_edges, 2]. Transpose if needed.
2. Forgetting activations: Check where the chosen layer or its supplied MLP applies nonlinearities.
3. Self-loops in hetero bipartite: Don't use add_self_loops=True when source and dest node types differ. Use skip connections instead.
4. NeighborLoader slicing: Only the first batch.batch_size nodes are your seed nodes. Slice predictions and labels accordingly.
5. Undirected graphs: If your graph is undirected, include edges in both directions in edge_index, or use T.ToUndirected().
6. Lazy init: Models with -1 input channels need one forward pass with torch.no_grad() before training to initialize parameters.
7. Global pooling for graph tasks: Use global_mean_pool(x, batch) (not manual reshape) to aggregate node features to graph-level.
8. num_neighbors alignment: Keep len(num_neighbors) equal to the number of GNN layers. More hops than layers wastes compute; fewer means wasted model capacity.

Choose splits before fitting features or model selection. Graph-level random splits can leak related molecules, patients, scaffolds, times, or sites; node-label masks define a transductive task unless unseen nodes/edges are excluded. Report the split unit, negative-edge universe, class balance, multiple seeds, and a task-appropriate baseline. A successful forward/backward pass is a mechanics check, not evidence of scientific generalization.