pytorch-geometric

# PyTorch Geometric — Graph Neural Networks PyTorch Geometric (PyG) is the standard library for deep learning on graphs. Where networkx handles graph algorithms (shortest path, centrality, community detection), PyG handles **learning on graphs**: training neural networks that operate directly on graph structure. The core insight: a GNN layer aggregates information from a node's neighbors, learns which neighbors matter, and produces new node representations — all differentiable, all trainable. ## Core Mental Model ``` A GRAPH has: • Nodes (vertices) — each has a feature vector • Edges (connections) — each optionally has attributes • Structure — which nodes connect to which A GNN LAYER does (per node): 1. GATHER messages from neighbors 2. AGGREGATE messages (sum / mean / max) 3. UPDATE own representation using aggregated + self Node v: h_v ← UPDATE( h_v, AGGREGATE( MESSAGE(h_u, e_uv) for u ∈ N(v) ) ) After k layers: each node "sees" its k-hop neighborhood. This is how local structure becomes global representation. PyG's DATA OBJECT: x → node feature matrix [num_nodes, num_features] edge_index → edge list (COO format) [2, num_edges] edge_attr → edge feature matrix [num_edges, num_edge_features] (optional) y → labels [num_nodes] or [num_graphs] (optional) ``` ### edge_index — The Key Format ``` Graph: 0 → 1, 0 → 2, 1 → 2 edge_index = tensor([[0, 0, 1], ← source nodes [1, 2, 2]]) ← target nodes Column i describes edge i: source = edge_index[0, i], target = edge_index[1, i] ⚠️ UNDIRECTED graph: store BOTH directions! 0 — 1 becomes 0→1 AND 1→0 → edge_index has 2× the edges Messages flow: source → target (default in MessagePassing) ``` ## Reference Documentation **PyG docs**: https://pytorch-geometric.readthedocs.io/en/latest/ **PyG tutorials**: https://pytorch-geometric.readthedocs.io/en/latest/tutorials.html **GitHub**: https://github.com/pyg-team/pytorch_geometric **Search patterns**: `Data`, `MessagePassing`, `GCNConv`, `global_mean_pool`, `Batch` ## Quick Reference ### Installation ```bash # Install PyTorch first (match your CUDA version) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # Then PyG — use the install finder at https://pytorch-geometric.readthedocs.io/en/latest/install.html pip install torch-geometric # Core extensions (required for many features): pip install torch-scatter torch-sparse torch-cluster torch-spline-conv ``` ### Standard Imports ```python import torch import torch.nn.functional as F from torch_geometric.data import Data, DataLoader, Batch from torch_geometric.nn import GCNConv, GATConv, GraphSAGEConv, GINConv from torch_geometric.nn import global_mean_pool, global_add_pool from torch_geometric.datasets import Planetoid, TUDataset ``` ### Basic Pattern — Build a Graph, Define a GNN, Train ```python import torch import torch.nn.functional as F from torch_geometric.data import Data from torch_geometric.nn import GCNConv # ─── 1. Build a graph ─── x = torch.tensor([[1.0, 0.0], # Node 0 features [0.0, 1.0], # Node 1 features [1.0, 1.0], # Node 2 features [0.0, 0.0]], # Node 3 features dtype=torch.float) # Edges: 0-1, 1-2, 2-3 (undirected → both directions) edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]], dtype=torch.long) y = torch.tensor([0, 0, 1, 1]) # Node labels (2 classes) data = Data(x=x, edge_index=edge_index, y=y) print(data) # Data(x=[4, 2], edge_index=[2, 6], y=[4]) # ─── 2. Define GNN model ─── class SimpleGCN(torch.nn.Module): def __init__(self, in_features, hidden, out_classes): super().__init__() self.conv1 = GCNConv(in_features, hidden) self.conv2 = GCNConv(hidden, out_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) x = self.conv2(x, edge_index) return x # [num_nodes, out_classes] — logits per node model = SimpleGCN(in_features=2, hidden=16, out_classes=2) out = model(data.x, data.edge_index) # Shape: [4, 2] ``` ## Critical Rules ### ✅ DO - **Make undirected graphs bidirectional in edge_index** — If edge 0→1 exists, include 1→0 too. Use `torch_geometric.utils.to_undirected()` to do this automatically. - **Keep edge_index as `torch.long` (int64)** — Always. Node feature tensors are float, edge_index must be long. - **Use `data.to(device)` to move entire graph** — Moves x, edge_index, edge_attr, y all at once. Don't move tensors individually. - **Use train_mask/val_mask/test_mask for node classification** — Standard transductive split. Masks are boolean tensors of shape [num_nodes]. - **Use DataLoader for graph classification** — It batches multiple graphs into one Batch object. Don't manually concatenate. - **Use `global_mean_pool` or `global_add_pool` before the final classifier in graph-level tasks** — Converts variable-size node matrices to fixed-size graph vectors. - **Add self-loops before GCN layers** — `GCNConv` adds them by default (`add_self_loops=True`). If you disabled them, node features don't propagate to themselves. - **Use `batch` argument in pooling** — `global_mean_pool(x, batch)` — `batch` tells the pooling which nodes belong to which graph in a Batch. ### ❌ DON'T - **Don't confuse edge_index shape** — It's `[2, E]`, NOT `[E, 2]`. Row 0 = sources, row 1 = targets. This is the #1 bug in PyG code. - **Don't use GCN on heterogeneous graphs** — GCNConv assumes homogeneous graphs (one node type, one edge type). Use `HeteroConv` or type-specific layers. - **Don't forget `model.eval()` and `torch.no_grad()` during inference** — Dropout and batch norm behave differently. - **Don't assume edge_index is sorted** — PyG doesn't guarantee edge ordering. Don't index into edge_attr assuming a specific edge order. - **Don't use standard PyTorch DataLoader** — Use `torch_geometric.data.DataLoader` which knows how to batch graphs. - **Don't stack node features across graphs manually** — `Batch.from_data_list()` handles this with correct edge_index offsetting. ## Anti-Patterns (NEVER) ```python import torch from torch_geometric.data import Data # ❌ BAD: edge_index transposed — [E, 2] instead of [2, E] edges = [(0,1), (1,2), (2,3)] edge_index = torch.tensor(edges, dtype=torch.long) # Shape: [3, 2] ← WRONG # GCNConv will silently produce garbage or crash. # ✅ GOOD: Transpose to [2, E] edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous() # Shape: [2, 3] ✓ # ───────────────────────────────────────────────────────────── # ❌ BAD: Directed edges for an undirected graph — messages flow one way only edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long) # Node 3 receives from 2, but node 0 never receives from 1. # GCN on this graph: nodes at the "end" of chains have rich representations, # nodes at the "start" stay at initialization. # ✅ GOOD: Add reverse edges from torch_geometric.utils import to_undirected edge_index = to_undirected(edge_index) # Now: [[0,1,1,2,2,3], [1,0,2,1,3,2]] # ───────────────────────────────────────────────────────────── # ❌ BAD: Moving tensors to different devices separately x = x.to('cuda') edge_index = edge_index.to('cuda') y = y.to('cuda') # Easy to forget one → runtime error # ✅ GOOD: Move the whole Data object data = data.to('cuda') # All attributes moved atomically # ───────────────────────────────────────────────────────────── # ❌ BAD: Standard DataLoader for graph datasets from torch.utils.data import DataLoader as TorchDataLoader loader = TorchDataLoader(dataset, batch_size=32) # Can't batch graphs! # ✅ GOOD: PyG DataLoader from torch_geometric.data import DataLoader loader = DataLoader(dataset, batch_size=32, shuffle=True) # Returns Batch objects — concatenated graphs with correct edge_index offsets ``` ## The Data Object ```python import torch from torch_geometric.data import Data from torch_geometric.utils import to_undirected # ─── Construct from scratch ─── data = Data( x=torch.randn(5, 16), # 5 nodes, 16 features each edge_index=torch.tensor([[0,1,2,3], [1,2,3,4]], dtype=torch.long), edge_attr=torch.randn(4, 8), # 4 edges, 8 features each y=torch.tensor([0, 1, 0, 1, 0]), # Node labels pos=torch.randn(5, 2), # Node positions (optional) ) # ─── Inspect ─── print(data) # Data(x=[5,16], edge_index=[2,4], ...) print(data.num_nodes) # 5 print(data.num_edges) # 4 print(data.num_node_features) # 16 print(data.is_undirected()) # True/False # ─── Make undirected ─── data.edge_index = to_undirected(data.edge_index) # edge_attr must also be duplicated if present: # data.edge_attr = torch.cat([data.edge_attr, data.edge_attr], dim=0) # ─── Build from edge list (e.g., from NetworkX or CSV) ─── import pandas as pd # Edge list: src, dst, weight edges_df = pd.DataFrame({ 'src': [0, 1, 2, 3], 'dst': [1, 2, 3, 0], 'weight': [0.5, 1.0, 0.3, 0.8] }) edge_index = torch.tensor([edges_df['src'].values, edges_df['dst'].values], dtype=torch.long) edge_index = to_undirected(edge_index) edge_attr = torch.tensor(edges_df['weight'].values, dtype=torch.float).unsqueeze(1) edge_attr = torch.cat([edge_attr, edge_attr], dim=0) # Mirror for undirected num_nodes = max(edges_df['src'].max(), edges_df['dst'].max()) + 1 x = torch.eye(num_nodes) # One-hot identity features if no attributes data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr) # ─── Add custom attributes ─── data.graph_label = torch.tensor([1]) # Graph-level label data.node_id = torch.arange(data.num_nodes) # Arbitrary metadata data.split = 'train' # String attributes fine too ``` ## MessagePassing Framework MessagePassing is the base class for ALL GNN layers in PyG. Understanding it = understanding how GNNs work. ```python import torch from torch_geometric.nn import MessagePassing import torch.nn.functional as F class CustomConv(MessagePassing): """ Custom GNN layer via MessagePassing. The propagate() call triggers this sequence: 1. message() — compute message for each edge (source → target) 2. aggregate() — combine messages arriving at each target node 3. update() — update each node's representation propagate(edge_index, x=x) routes: • x_j → source node features (j = source index) • x_i → target node features (i = target index) Subscript _i = target, _j = source. Always. """ def __init__(self, in_channels, out_channels): super().__init__(aggr='add') # Aggregation: 'add', 'mean', 'max' self.lin = torch.nn.Linear(in_channels, out_channels) def forward(self, x, edge_index): # Transform features BEFORE propagation (more efficient) x = self.lin(x) # Propagate: runs message() → aggregate() → update() return self.propagate(edge_index, x=x) def message(self, x_j): """ x_j: source node features for each edge. Shape: [num_edges, out_channels] Return: message to send along each edge. """ return x_j # Simplest: just pass source features through def update(self, aggr_out): """ aggr_out: aggregated messages per target node. Shape: [num_nodes, out_channels] Return: updated node representation. """ return aggr_out # Simplest: use aggregation directly # ─── Attention-weighted custom layer ─── class AttentionConv(MessagePassing): """Messages weighted by learned attention scores (simplified GAT).""" def __init__(self, in_channels, out_channels): super().__init__(aggr='add') self.lin = torch.nn.Linear(in_channels, out_channels) self.att = torch.nn.Parameter(torch.Tensor(1, out_channels)) torch.nn.init.xavier_uniform_(self.att.unsqueeze(0)) def forward(self, x, edge_index): x = self.lin(x) return self.propagate(edge_index, x=x) def message(self, x_i, x_j): # x_i = target features, x_j = source features # Attention score: how much should target i attend to source j? alpha = (x_i * self.att).sum(dim=-1) + (x_j * self.att).sum(dim=-1) alpha = F.leaky_relu(alpha, 0.2) # Note: full GAT uses softmax over neighbors — see GATConv for production version return x_j * alpha.unsqueeze(-1) # ─── Usage ─── # conv = CustomConv(16, 32) # out = conv(data.x, data.edge_index) # [num_nodes, 32] ``` ## Standard Layers — When to Use Which ```python from torch_geometric.nn import GCNConv, GATConv, GraphSAGEConv, GINConv import torch.nn as nn # ─── GCNConv: Graph Convolutional Network (Kipf & Welling, 2017) ─── # Averages neighbor features (with degree normalization). # Fast, simple baseline. No attention, no edge features. # USE: citation networks, social networks, when speed matters. conv_gcn = GCNConv(in_channels=16, out_channels=32) # out = conv_gcn(x, edge_index) # ─── GATConv: Graph Attention Network (Veličković et al., 2018) ─── # Learns attention weights — some neighbors matter more than others. # Slower than GCN but usually better accuracy. # USE: when neighbor importance varies, heterogeneous neighborhood structure.