From Graph Math to PyTorch Implementation

Unlike images (regular 2D grids) or text (1D sequences), graph structures represent relational data with arbitrary topologies. Graph Neural Networks (GNNs) operate directly on non-Euclidean data by combining graph topology with node feature transformations.

This guide covers graph data representations in code, the mechanics of message passing, building custom GNN layers from scratch, and implementing real-world node and graph classification pipelines using PyTorch and PyTorch Geometric (PyG).

1. Graph Fundamentals & Tensor Representation

A graph is represented as $G = (V, E, \mathbf{X})$, where:

• $V$ is the set of $\vert{}V\vert{}$ vertices (nodes).
• $E$ is the set of $\vert{}E\vert{}$ edges between vertices.
• $\mathbf{X} \in \mathbb{R}^{\vert{}V\vert{} \times d}$ is the node feature matrix, where each node has a feature vector of dimension $d$.

To store edge topologies efficiently, sparse formats are preferred over dense Adjacency Matrices ($\mathbf{A} \in \mathbb{R}^{\vert{}V\vert{} \times \vert{}V\vert{}}$). PyTorch Geometric uses the COO (Coordinate Format) sparse representation: a 2D tensor $\mathbf{E} \in \mathbb{Z}^{2 \times \vert{}E\vert{}}$ storing target and source node pairs.

Dense Adjacency Matrix (A)         COO Edge Index Format (E)
    Node 0  1  2  3                   
0  [  0, 1, 1, 0 ]                 Row 0 (Sources): [0, 0, 1, 2, 3]
1  [  1, 0, 0, 1 ]                 Row 1 (Targets): [1, 2, 0, 3, 1]
2  [  1, 0, 0, 0 ]                 
3  [  0, 1, 0, 0 ]                 Shape: [2, E]

Step 1: Creating Graph Data Structures in PyTorch Geometric

Python

import torch
from torch_geometric.data import Data

# 1. Define Node Features (4 nodes, each with 3 feature dimensions)
# Nodes: 0, 1, 2, 3
x = torch.tensor([
    [-1.0,  0.5,  1.2],  # Node 0
    [ 0.0,  1.0, -0.5],  # Node 1
    [ 2.1, -1.1,  0.0],  # Node 2
    [ 0.4,  0.0,  0.8]   # Node 3
], dtype=torch.float)

# 2. Define Graph Topology in COO format [2, num_edges]
# Directed edges: (0->1), (1->0), (0->2), (2->0), (1->3), (3->1)
edge_index = torch.tensor([
    [0, 1, 0, 2, 1, 3],  # Source nodes
    [1, 0, 2, 0, 3, 1]   # Target nodes
], dtype=torch.long)

# 3. Define Node Labels (e.g., binary node classification targets)
y = torch.tensor([0, 1, 0, 1], dtype=torch.long)

# 4. Construct PyG Data Object
graph = Data(x=x, edge_index=edge_index, y=y)

print("Graph Summary:")
print(f"Number of nodes: {graph.num_nodes}")
print(f"Number of edges: {graph.num_edges}")
print(f"Node feature matrix shape: {graph.x.shape}")
print(f"Contains isolated nodes: {graph.has_isolated_nodes()}")
print(f"Is undirected: {graph.is_undirected()}")

2. The Message Passing Framework

The unified abstraction behind GNN architectures (GCN, GraphSAGE, GAT) is Message Passing. At layer $k$, the update of node $v$’s hidden state $\mathbf{h}_v^{(k)}$ proceeds in three steps:

$$\mathbf{m}_{u \to v}^{(k)} = \text{MSG}^{(k)}\left(\mathbf{h}_u^{(k-1)}, \mathbf{h}_v^{(k-1)}, \mathbf{e}_{u \to v}\right)$$

$$\mathbf{m}_v^{(k)} = \mathbf{\square}_{u \in \mathcal{N}(v)} \, \mathbf{m}_{u \to v}^{(k)}$$

$$\mathbf{h}_v^{(k)} = \text{UPDATE}^{(k)}\left(\mathbf{h}_v^{(k-1)}, \mathbf{m}_v^{(k)}\right)$$

Where:

• $\text{MSG}(\cdot)$ generates a message from neighbor node $u$ to target node $v$.
• $\mathbf{\square}$ is a permutation-invariant aggregation operator ($\sum$, $\text{Mean}$, or $\text{Max}$).
• $\text{UPDATE}(\cdot)$ combines the aggregated message with node $v$’s previous representation.

Neighborhood Aggregation for Target Node v:

     (Neighbor u1) [h_u1] ──┐
                            ├─► [ MSG ] ─┐
     (Neighbor u2) [h_u2] ──┤            │
                            ├─► [ MSG ] ─┼─► [ AGGREGATE (Σ) ] ─► [ UPDATE (MLP) ] ─► h_v^(k)
     (Neighbor u3) [h_u3] ──┘            │
                            └─► [ MSG ] ─┘

3. Building a GNN Layer from Scratch

To understand the core operations without high-level PyG abstractions, we can build a basic Graph Convolutional Network (GCN) layer using raw PyTorch.

The mathematical formulation for a symmetric-normalized GCN layer (Kipf & Welling) is:

$$\mathbf{H}^{(k+1)} = \sigma \left( \tilde{\mathbf{D}}^{-\frac{1}{2}} \tilde{\mathbf{A}} \tilde{\mathbf{D}}^{-\frac{1}{2}} \mathbf{H}^{(k)} \mathbf{W}^{(k)} \right)$$

Where $\tilde{\mathbf{A}} = \mathbf{A} + \mathbf{I}_N$ (adjacency matrix with added self-loops) and $\tilde{\mathbf{D}}_{ii} = \sum_j \tilde{\mathbf{A}}_{ij}$ (degree matrix).

Step 2: Custom GCN Layer Implementation

Python

import torch
import torch.nn as nn

class CustomGCNLayer(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super(CustomGCNLayer, self).__init__()
        # Trainable weight transformation matrix W
        self.linear = nn.Linear(in_features, out_features, bias=False)
        
    def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
        num_nodes = x.size(0)
        
        # 1. Add Self-Loops: A_tilde = A + I
        loop_index = torch.arange(0, num_nodes, dtype=torch.long, device=edge_index.device)
        loop_index = loop_index.unsqueeze(0).repeat(2, 1)
        edge_index_tilde = torch.cat([edge_index, loop_index], dim=1)
        
        # 2. Compute Degree Matrix D_tilde
        src, dst = edge_index_tilde[0], edge_index_tilde[1]
        deg = torch.zeros(num_nodes, dtype=torch.float, device=x.device)
        deg.index_add_(0, src, torch.ones_like(src, dtype=torch.float))
        
        # 3. Compute Degree Normalization Coefficients: D^(-1/2)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0.0
        
        # Compute edge weights: norm = deg_inv_sqrt[src] * deg_inv_sqrt[dst]
        norm = deg_inv_sqrt[src] * deg_inv_sqrt[dst]
        
        # 4. Feature Transformation: H * W
        h_transformed = self.linear(x)
        
        # 5. Message Passing (Scatter Aggregation)
        # Aggregating normalized neighbor features into target nodes
        out = torch.zeros_like(h_transformed)
        messages = h_transformed[src] * norm.unsqueeze(-1)
        out.index_add_(0, dst, messages)
        
        return out

# Verification
custom_layer = CustomGCNLayer(in_features=3, out_features=4)
output_features = custom_layer(graph.x, graph.edge_index)

print(f"Input feature shape:  {graph.x.shape}")
print(f"Output feature shape: {output_features.shape}")

4. Node Classification Pipeline with PyTorch Geometric

Now we build a production-grade Node Classification model using standard PyG components on the Cora citation network dataset.

Step 3: Dataset Preparation, Model Definition, and Training

Python

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

# 1. Load the Cora Dataset
dataset = Planetoid(root='/tmp/Cora', name='Cora')
cora_data = dataset[0]

print(f"Dataset: {dataset.name}")
print(f"Nodes: {cora_data.num_nodes}, Edges: {cora_data.num_edges}")
print(f"Features per node: {dataset.num_features}")
print(f"Number of target classes: {dataset.num_classes}")

# 2. Define a 2-Layer GCN Architecture
class GCNNodeClassifier(nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int, out_channels: int, dropout: float = 0.5):
        super(GCNNodeClassifier, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)
        self.dropout = dropout

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
        # First Message Passing Layer + Non-linearity + Dropout
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        
        # Second Message Passing Layer (Logits output)
        x = self.conv2(x, edge_index)
        return x

# 3. Instantiate Model, Loss, and Optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCNNodeClassifier(
    in_channels=dataset.num_features,
    hidden_channels=16,
    out_channels=dataset.num_classes,
    dropout=0.5
).to(device)

data = cora_data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()

# 4. Define Train and Evaluation Functions
def train() -> float:
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    
    # Compute loss ONLY on training nodes using training mask
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def evaluate(mask: torch.Tensor) -> float:
    model.eval()
    out = model(data.x, data.edge_index)
    preds = out.argmax(dim=-1)
    correct = (preds[mask] == data.y[mask]).sum()
    acc = int(correct) / int(mask.sum())
    return acc

# 5. Training Loop
print("\nStarting Training...")
for epoch in range(1, 101):
    loss = train()
    if epoch % 20 == 0 or epoch == 1:
        val_acc = evaluate(data.val_mask)
        test_acc = evaluate(data.test_mask)
        print(f"Epoch: {epoch:03d} | Train Loss: {loss:.4f} | Val Acc: {val_acc:.4f} | Test Acc: {test_acc:.4f}")

5. Graph Classification & Global Readout Operations

While node classification computes embeddings for individual nodes ($\mathbf{h}_v$), Graph Classification maps an entire graph topology to a single graph-level vector ($\mathbf{h}_G$).

This requires a Global Readout (Pooling) step to aggregate node features across the entire graph dimension:

$$\mathbf{h}_G = \text{READOUT}\left( \{\mathbf{h}_v^{(L)} \mid v \in V\} \right) = \sum_{v \in V} \mathbf{h}_v^{(L)} \quad \text{or} \quad \frac{1}{\vert{}V\vert{}}\sum_{v \in V} \mathbf{h}_v^{(L)}$$

Graph-Level Pooling Architecture:

 Node 0 [h_0] ──┐
 Node 1 [h_1] ──┼─► [ GLOBAL READOUT ] ──► [ Graph Embedding h_G ] ──► [ Linear Classifier ]
 Node 2 [h_2] ──┘   (Mean / Sum Pool)                                 (Graph Output Class)

Step 4: Graph Classification Pipeline on MUTAG Dataset

Python

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool

# 1. Load MUTAG Graph Classification Dataset
dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG')

# Split dataset into train/test sets
torch.manual_seed(42)
dataset = dataset.shuffle()
train_dataset = dataset[:150]
test_dataset = dataset[150:]

# DataLoader automatically constructs disjoint mini-batch graphs (batch indices)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 2. Define Graph Classifier Module
class GraphClassifier(nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int, out_channels: int):
        super(GraphClassifier, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)
        
        # Final classification linear head
        self.classifier = nn.Linear(hidden_channels, out_channels)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, batch: torch.Tensor) -> torch.Tensor:
        # 1. Obtain Node Embeddings via GNN layers
        h = F.relu(self.conv1(x, edge_index))
        h = F.relu(self.conv2(h, edge_index))
        h = self.conv3(h, edge_index)

        # 2. Readout Layer: Aggregate node features into graph-level representations
        # batch tensor maps each node to its respective graph index in the mini-batch
        h_graph = global_mean_pool(h, batch)  # Shape: [batch_size, hidden_channels]

        # 3. Apply Classifier Head
        h_graph = F.dropout(h_graph, p=0.5, training=self.training)
        out = self.classifier(h_graph)
        return out

# 3. Training setup
model = GraphClassifier(
    in_channels=dataset.num_node_features,
    hidden_channels=64,
    out_channels=dataset.num_classes
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

def train_graph_classifier():
    model.train()
    total_loss = 0
    for data in train_loader:
        optimizer.zero_grad()
        # Pass node features, edges, and batch index vector
        out = model(data.x, data.edge_index, data.batch)
        loss = criterion(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
    return total_loss / len(train_loader.dataset)

@torch.no_grad()
def test_graph_classifier(loader):
    model.eval()
    correct = 0
    for data in loader:
        out = model(data.x, data.edge_index, data.batch)
        pred = out.argmax(dim=-1)
        correct += int((pred == data.y).sum())
    return correct / len(loader.dataset)

# 4. Execution Loop
print("Training Graph Classifier on MUTAG...")
for epoch in range(1, 51):
    loss = train_graph_classifier()
    if epoch % 10 == 0 or epoch == 1:
        train_acc = test_graph_classifier(train_loader)
        test_acc = test_graph_classifier(test_loader)
        print(f"Epoch: {epoch:02d} | Train Loss: {loss:.4f} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

6. Model Comparison & Architecture Choice Guide

Selecting the appropriate GNN operator depends on graph density, computational constraints, and task structure:

Architectural Operator	Aggregation Function	Computational Bottleneck	Best Used For
GCN (Graph Convolutional)	Degree-normalized mean summation	Memory bandwidth (Sparse matrix ops)	Homophilic networks (Citation networks, social graphs)
GraphSAGE (Sample & Aggregate)	Uniform neighborhood sampling + Aggregation ($\text{Mean}/\text{LSTM}$)	Sampling overhead & I/O transfer	Large-scale dynamic graphs (E-commerce, recommendation systems)
GAT (Graph Attention Network)	Anisotropic weighted sum (Self-attention weights)	Compute intensity (Dynamic edge weight calculation)	Graphs with varying edge significance (Molecular structures, protein interactions)