Graph Neural Networks (GNNs) extend deep learning to graph-structured data — social networks, molecular structures, knowledge graphs, and fraud detection networks. When relationships between entities matter as much as the entities themselves, GNNs outperform traditional ML models. This guide builds working GNNs using PyTorch Geometric.
Why Graphs?
Many real-world problems are naturally graph-structured: fraud rings (accounts connected by shared devices), drug discovery (atoms connected by chemical bonds), recommendation systems (users connected to items), and social networks (people connected by relationships). Traditional ML models treat each sample independently — they cannot leverage the information encoded in connections. GNNs propagate information across edges so each node’s representation is informed by its neighbourhood.
Installing PyTorch Geometric
pip install torch torchvision torchaudio
pip install torch_geometric
pip install torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.1.0+cpu.html
import torch
import torch_geometric
from torch_geometric.data import Data
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv, SAGEConv, GATConv
import torch.nn.functional as F
print(torch_geometric.__version__)
Graph Data Structure
import torch
from torch_geometric.data import Data
# Build a simple graph manually
# 4 nodes, each with 3 features
x = torch.tensor([
[1, 0, 1], # node 0
[0, 1, 0], # node 1
[1, 1, 0], # node 2
[0, 0, 1], # node 3
], dtype=torch.float)
# Edges: (0-1), (1-2), (2-3), (0-3) — bidirectional
edge_index = torch.tensor([
[0, 1, 1, 2, 2, 3, 3, 0], # source nodes
[1, 0, 2, 1, 3, 2, 0, 3], # target nodes
], dtype=torch.long)
# Node labels for classification
y = torch.tensor([0, 1, 0, 1], dtype=torch.long)
graph = Data(x=x, edge_index=edge_index, y=y)
print(graph)
print(f'Nodes: {graph.num_nodes} | Edges: {graph.num_edges}')
print(f'Node features: {graph.num_node_features}')
Graph Convolutional Network (GCN)
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
# Cora citation network: 2708 papers, 5429 citations, 7 classes
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
print(f'Nodes: {data.num_nodes} | Edges: {data.num_edges}')
print(f'Features: {data.num_features} | Classes: {dataset.num_classes}')
print(f'Train mask: {data.train_mask.sum()} nodes')
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)
self.dropout = torch.nn.Dropout(0.5)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index)
x = F.relu(x)
x = self.dropout(x)
x = self.conv2(x, edge_index)
return x # raw logits
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN(dataset.num_features, 64, dataset.num_classes).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
def train():
model.train()
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()
return loss.item()
def evaluate(mask):
model.eval()
with torch.no_grad():
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1)
acc = (pred[mask] == data.y[mask]).float().mean()
return acc.item()
best_val_acc = 0
for epoch in range(200):
loss = train()
val_acc = evaluate(data.val_mask)
if val_acc > best_val_acc:
best_val_acc = val_acc
torch.save(model.state_dict(), 'best_gcn.pt')
if epoch % 20 == 0:
print(f'Epoch {epoch:3d} | Loss: {loss:.4f} | Val Acc: {val_acc:.4f}')
model.load_state_dict(torch.load('best_gcn.pt'))
test_acc = evaluate(data.test_mask)
print(f'Test Accuracy: {test_acc:.4f}')
Graph Attention Network (GAT)
from torch_geometric.nn import GATConv
class GAT(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels,
heads=8, dropout=0.6):
super().__init__()
self.dropout = dropout
self.conv1 = GATConv(in_channels, hidden_channels,
heads=heads, dropout=dropout)
self.conv2 = GATConv(hidden_channels * heads, out_channels,
heads=1, concat=False, dropout=dropout)
def forward(self, x, edge_index):
x = F.dropout(x, p=self.dropout, training=self.training)
x = F.elu(self.conv1(x, edge_index))
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv2(x, edge_index)
return x
gat = GAT(dataset.num_features, 8, dataset.num_classes).to(device)
# GAT uses attention to weight neighbour contributions — better for
# heterogeneous graphs where not all neighbours are equally relevant
Link Prediction (Fraud Detection)
from torch_geometric.utils import negative_sampling
from torch_geometric.nn import SAGEConv
class GraphSAGE(torch.nn.Module):
def __init__(self, in_channels, hidden_channels):
super().__init__()
self.conv1 = SAGEConv(in_channels, hidden_channels)
self.conv2 = SAGEConv(hidden_channels, hidden_channels)
def encode(self, x, edge_index):
x = self.conv1(x, edge_index).relu()
return self.conv2(x, edge_index)
def decode(self, z, edge_index):
# Dot product of node embeddings predicts link existence
return (z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)
def forward(self, x, edge_index, neg_edge_index):
z = self.encode(x, edge_index)
pos_scores = self.decode(z, edge_index)
neg_scores = self.decode(z, neg_edge_index)
return pos_scores, neg_scores
sage = GraphSAGE(dataset.num_features, 64).to(device)
optimizer = torch.optim.Adam(sage.parameters(), lr=0.01)
for epoch in range(100):
sage.train()
optimizer.zero_grad()
neg_edge = negative_sampling(data.edge_index,
num_nodes=data.num_nodes,
num_neg_samples=data.edge_index.size(1))
pos_scores, neg_scores = sage(data.x, data.edge_index, neg_edge.to(device))
labels = torch.cat([torch.ones(pos_scores.size(0)),
torch.zeros(neg_scores.size(0))]).to(device)
scores = torch.cat([pos_scores, neg_scores])
loss = F.binary_cross_entropy_with_logits(scores, labels)
loss.backward()
optimizer.step()
if epoch % 20 == 0:
print(f'Epoch {epoch}: Loss = {loss.item():.4f}')
Conclusion
GNNs unlock a new class of problems that tabular models fundamentally cannot solve: anything where relationships between entities carry predictive signal. Node classification, link prediction, and graph classification each have established GNN architectures. Start with GCN for homogeneous graphs, GAT when neighbour importance varies, and GraphSAGE for large graphs that require mini-batch training. The PyTorch Geometric ecosystem makes all of these accessible with clean, composable APIs.



