gcn_cora.py — node classification with GCN
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.datasets import Planetoid
# Load Cora citation network
dataset = Planetoid(root='data/', name='Cora')
data = dataset[0]
class GCN(torch.nn.Module):
def __init__(self, in_channels, hidden, out_channels):
super().__init__()
self.conv1 = GCNConv(in_channels, hidden)
self.conv2 = GCNConv(hidden, out_channels)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=0.5, training=self.training)
x = self.conv2(x, edge_index)
return F.log_softmax(x, dim=1)
model = GCN(dataset.num_features, 64, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
# Training loop
for epoch in range(200):
model.train()
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
# Evaluate
model.eval()
pred = model(data.x, data.edge_index).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}") # ~81.5%
gat.py — Graph Attention Network
from torch_geometric.nn import GATConv
class GAT(torch.nn.Module):
"""Graph Attention Network - learns WHICH neighbors matter."""
def __init__(self, in_channels, hidden, out_channels, heads=8):
super().__init__()
self.conv1 = GATConv(in_channels, hidden, heads=heads, dropout=0.6)
self.conv2 = GATConv(hidden * heads, out_channels, heads=1,
concat=False, dropout=0.6)
def forward(self, x, edge_index):
x = F.dropout(x, p=0.6, training=self.training)
x = F.elu(self.conv1(x, edge_index))
x = F.dropout(x, p=0.6, training=self.training)
x = self.conv2(x, edge_index)
return F.log_softmax(x, dim=1)
# GAT typically hits ~83.0% on Cora (vs ~81.5% for GCN)
# The attention weights are interpretable - you can visualize them
graphsage_minibatch.py — scalable inductive learning
from torch_geometric.nn import SAGEConv
from torch_geometric.loader import NeighborLoader
class GraphSAGE(torch.nn.Module):
"""Inductive learning - works on unseen nodes/graphs."""
def __init__(self, in_channels, hidden, out_channels):
super().__init__()
self.conv1 = SAGEConv(in_channels, hidden)
self.conv2 = SAGEConv(hidden, out_channels)
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
# Mini-batch training for large graphs (millions of nodes)
loader = NeighborLoader(
data,
num_neighbors=[25, 10], # Sample 25 1-hop, 10 2-hop neighbors
batch_size=1024,
input_nodes=data.train_mask,
)
for batch in loader:
out = model(batch.x, batch.edge_index)
loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
loss.backward()
optimizer.step()
gin_molecules.py — graph-level classification with GIN
from torch_geometric.nn import GINConv, global_add_pool
from torch_geometric.datasets import MoleculeNet
class GIN(torch.nn.Module):
"""Graph Isomorphism Network - maximally expressive message passing."""
def __init__(self, in_channels, hidden, out_channels, num_layers=5):
super().__init__()
self.convs = torch.nn.ModuleList()
self.bns = torch.nn.ModuleList()
for i in range(num_layers):
dim_in = in_channels if i == 0 else hidden
mlp = torch.nn.Sequential(
torch.nn.Linear(dim_in, hidden),
torch.nn.ReLU(),
torch.nn.Linear(hidden, hidden),
)
self.convs.append(GINConv(mlp))
self.bns.append(torch.nn.BatchNorm1d(hidden))
self.classifier = torch.nn.Linear(hidden, out_channels)
def forward(self, x, edge_index, batch):
for conv, bn in zip(self.convs, self.bns):
x = F.relu(bn(conv(x, edge_index)))
# Global pooling: graph-level readout
x = global_add_pool(x, batch)
return self.classifier(x)
# Molecular property prediction
dataset = MoleculeNet(root='data/', name='ogbg-molpcba')