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

# 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

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

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 edgeindex as torch.long (int64) — Always. Node feature tensors are float, edgeindex must be long.
• Use data.to(device) to move entire graph — Moves x, edgeindex, edgeattr, y all at once. Don't move tensors individually.
• Use trainmask/valmask/testmask for node classification — Standard transductive split. Masks are boolean tensors of shape [numnodes].
• 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 edgeindex is sorted — PyG doesn't guarantee edge ordering. Don't index into edgeattr 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)

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

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.

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

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.
conv_gat = GATConv(in_channels=16, out_channels=32, heads=8, concat=True)
# out shape: [num_nodes, heads * out_channels] if concat=True
# out = conv_gat(x, edge_index)

# ─── GraphSAGEConv: Inductive Representation Learning (Hamilton et al., 2017) ───
# Samples and aggregates neighbor features. Designed for INDUCTIVE setting
# (generalizes to unseen nodes at test time, e.g., new users joining a network).
# USE: large dynamic graphs, inductive node classification.
conv_sage = GraphSAGEConv(in_channels=16, out_channels=32, aggr='mean')
# out = conv_sage(x, edge_index)

# ─── GINConv: Graph Isomorphism Network (Xu et al., 2019) ───
# Theoretically most powerful among message-passing GNNs.
# Uses a learnable MLP to combine self + aggregated neighbors.
# USE: graph classification, when distinguishing graph structures matters.
mlp = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 32))
conv_gin = GINConv(nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 32)))
# out = conv_gin(x, edge_index)

# ─── LAYER SELECTION GUIDE ───
# Task                          | Recommended
# ──────────────────────────────┼─────────────────────
# Citation network (transductive)| GCN or GAT
# Social network (inductive)    | GraphSAGE
# Molecular property prediction | GIN or GAT
# Knowledge graph               | GAT (edge types matter)
# Quick baseline                | GCN (fastest)
# Best accuracy (small graph)   | GAT (multi-head)
# Best graph classification     | GIN (most expressive)

Node Classification Pipeline

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv, GATConv

# ─── Load dataset: Cora citation network ───
# 2708 papers (nodes), 5429 citations (edges), 7 classes
# Each node has a 1433-dim bag-of-words feature vector
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data    = dataset[0]   # Single graph

print(f"Nodes: {data.num_nodes}, Edges: {data.num_edges}, "
      f"Features: {data.num_node_features}, Classes: {dataset.num_classes}")
print(f"Train: {data.train_mask.sum()}, Val: {data.val_mask.sum()}, "
      f"Test: {data.test_mask.sum()}")

# ─── Model ───
class GATNodeClassifier(torch.nn.Module):
    def __init__(self, in_features, hidden, out_classes, heads=8, dropout=0.6):
        super().__init__

…

## Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

- **Author:** [tondevrel](https://github.com/tondevrel)
- **Source:** [tondevrel/scientific-agent-skills](https://github.com/tondevrel/scientific-agent-skills)
- **License:** MIT

Install and usage instructions live in the source repository linked above.