PyTorch Geometric is a library built on PyTorch for developing and training Graph Neural Networks (GNNs). Apply this skill for deep learning on graphs and irregular structures, including mini-batch processing, multi-GPU training, and geometric deep learning applications.
When to Use This Skill
This skill should be used when working with:
Graph-based machine learning: Node classification, graph classification, link prediction
Molecular property prediction: Drug discovery, chemical property prediction
Social network analysis: Community detection, influence prediction
Citation networks: Paper classification, recommendation systems
3D geometric data: Point clouds, meshes, molecular structures
Heterogeneous graphs: Multi-type nodes and edges (e.g., knowledge graphs)
Large-scale graph learning: Neighbor sampling, distributed training
Quick Start
Installation
uv pip install torch_geometric
For additional dependencies (sparse operations, clustering):
Check references/datasets_reference.md for a comprehensive list.
Creating Custom Datasets
For datasets that fit in memory, inherit from InMemoryDataset:
from torch_geometric.data import InMemoryDataset, Data
import torch
classMyOwnDataset(InMemoryDataset):
def__init__(self, root, transform=None, pre_transform=None):
super().__init__(root, transform, pre_transform)
self.load(self.processed_paths[0])
@propertydefraw_file_names(self):
return ['my_data.csv'] # Files needed in raw_dir @propertydefprocessed_file_names(self):
return ['data.pt'] # Files in processed_dirdefdownload(self):
# Download raw data to self.raw_dirpassdefprocess(self):
# Read data, create Data objects
data_list = []
# Example: Create a simple graph
edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
x = torch.randn(2, 16)
y = torch.tensor([0], dtype=torch.long)
data = Data(x=x, edge_index=edge_index, y=y)
data_list.append(data)
# Apply pre_filter and pre_transformifself.pre_filter isnotNone:
data_list = [d for d in data_list ifself.pre_filter(d)]
ifself.pre_transform isnotNone:
data_list = [self.pre_transform(d) for d in data_list]
# Save processed dataself.save(data_list, self.processed_paths[0])
For large datasets that don't fit in memory, inherit from Dataset and implement len() and get(idx).
Loading Graphs from CSV
import pandas as pd
import torch
from torch_geometric.data import HeteroData
# Load nodes
nodes_df = pd.read_csv('nodes.csv')
x = torch.tensor(nodes_df[['feat1', 'feat2']].values, dtype=torch.float)
# Load edges
edges_df = pd.read_csv('edges.csv')
edge_index = torch.tensor([edges_df['source'].values,
edges_df['target'].values], dtype=torch.long)
data = Data(x=x, edge_index=edge_index)
Training Workflows
Node Classification (Single Graph)
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
# Load dataset
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
# Create model
model = GCN(dataset.num_features, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
# Training
model.train()
for epoch inrange(200):
optimizer.zero_grad()
out = model(data)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f'Epoch {epoch}, Loss: {loss.item():.4f}')
# Evaluation
model.eval()
pred = model(data).argmax(dim=1)
correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
acc = int(correct) / int(data.test_mask.sum())
print(f'Test Accuracy: {acc:.4f}')
Graph Classification (Multiple Graphs)
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import global_mean_pool
classGraphClassifier(torch.nn.Module):
def__init__(self, num_features, num_classes):
super().__init__()
self.conv1 = GCNConv(num_features, 64)
self.conv2 = GCNConv(64, 64)
self.lin = torch.nn.Linear(64, num_classes)
defforward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = self.conv1(x, edge_index)
x = F.relu(x)
x = self.conv2(x, edge_index)
x = F.relu(x)
# Global pooling (aggregate node features to graph-level)
x = global_mean_pool(x, batch)
x = self.lin(x)
return F.log_softmax(x, dim=1)
# Load dataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
loader = DataLoader(dataset, batch_size=32, shuffle=True)
model = GraphClassifier(dataset.num_features, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# Training
model.train()
for epoch inrange(100):
total_loss = 0for batch in loader:
optimizer.zero_grad()
out = model(batch)
loss = F.nll_loss(out, batch.y)
loss.backward()
optimizer.step()
total_loss += loss.item()
if epoch % 10 == 0:
print(f'Epoch {epoch}, Loss: {total_loss / len(loader):.4f}')
Large-Scale Graphs with Neighbor Sampling
For large graphs, use NeighborLoader to sample subgraphs:
from torch_geometric.loader import NeighborLoader
# Create a neighbor sampler
train_loader = NeighborLoader(
data,
num_neighbors=[25, 10], # Sample 25 neighbors for 1st hop, 10 for 2nd hop
batch_size=128,
input_nodes=data.train_mask,
)
# Training
model.train()
for batch in train_loader:
optimizer.zero_grad()
out = model(batch)
# Only compute loss on seed nodes (first batch_size nodes)
loss = F.nll_loss(out[:batch.batch_size], batch.y[:batch.batch_size])
loss.backward()
optimizer.step()
Important:
Output subgraphs are directed
Node indices are relabeled (0 to batch.num_nodes - 1)
Only use seed node predictions for loss computation
Sampling beyond 2-3 hops is generally not feasible
Advanced Features
Heterogeneous Graphs
For graphs with multiple node and edge types, use HeteroData:
from torch_geometric.data import HeteroData
data = HeteroData()
# Add node features for different types
data['paper'].x = torch.randn(100, 128) # 100 papers with 128 features
data['author'].x = torch.randn(200, 64) # 200 authors with 64 features# Add edges for different types (source_type, edge_type, target_type)
data['author', 'writes', 'paper'].edge_index = torch.randint(0, 200, (2, 500))
data['paper', 'cites', 'paper'].edge_index = torch.randint(0, 100, (2, 300))
print(data)
Convert homogeneous models to heterogeneous:
from torch_geometric.nn import to_hetero
# Define homogeneous model
model = GNN(...)
# Convert to heterogeneous
model = to_hetero(model, data.metadata(), aggr='sum')
# Use as normal
out = model(data.x_dict, data.edge_index_dict)
Or use HeteroConv for custom edge-type-specific operations:
