Graph Neural Networks: Architectures, Training, and Production Deployment

Beyond Grids and Sequences

Deep learning conquered images with convolutions and language with transformers. Both domains assume a regular structure — pixels in a grid, tokens in a sequence. But the world's most interesting data lives in graphs: social networks, molecular structures, citation networks, knowledge graphs, and supply chains. These have no fixed grid, no sequential order. Each graph has a different number of nodes with a different arrangement of edges.

Graph Neural Networks (GNNs)¹ extend deep learning to this irregular domain. Instead of sliding a kernel over a grid, GNNs propagate information along edges, letting each node aggregate features from its neighbours. The result is a representation that captures both the node's own properties and the structure of its local neighbourhood — exactly what you need for tasks like predicting protein interactions, ranking documents in a knowledge graph, or detecting fraud in a transaction network.

This article surveys the major GNN architectures, shows you how to train them with PyTorch Geometric, and covers what it takes to deploy them to production.

The Message-Passing Paradigm

Every GNN is built on the same core operation: message passing. For each node, the GNN collects messages from its neighbours, aggregates them (sum, mean, or attention-weighted), and updates the node's representation. After one round, each node knows about its direct neighbours. After k rounds, each node knows about nodes up to k hops away.

h_v^(0) = x_v                          (initial node features)

h_v^(k) = UPDATE(
    h_v^(k-1),
    AGGREGATE({ h_u^(k-1) for every neighbour u of v })
)

This is deceptively simple. The choice of aggregate and update functions determines the entire GNN's behaviour, and different architectures make different trade-offs between expressiveness, computational cost, and over-smoothing.

GNN Architectures Compared

Architecture	Aggregation	Key Innovation	Best For	Limitations
GCN (Kipf & Welling, 2017)	Normalised mean	Symmetric normalisation of adjacency matrix	Citation networks, node classification	Treats all neighbours equally
GraphSAGE (Hamilton et al., 2017)	Mean / LSTM / Pool	Inductive — learns for unseen nodes; supports sampling	Large graphs, evolving graphs	Higher memory than GCN with LSTM aggregator
GAT (Veličković et al., 2018)	Attention-weighted sum	Self-attention over neighbours — learns which neighbours matter	Heterogeneous graphs, relationship-rich domains	O(E) attention computation
GIN (Xu et al., 2019)	Sum + MLP	Maximally expressive under WL-test	Graph classification, isomorphism detection	Deeper networks still over-smooth
Graph Transformer (Dwivedi & Bresson, 2021)	Full attention	Node positional encoding + transformer attention	Large graphs with long-range dependencies	O(V²) attention — does not scale beyond ~5k nodes

GCN is where most people start. It is simple, fast, and works well for homophilic graphs — where connected nodes tend to be similar (citation networks, social networks). GAT is the go-to when edges carry different meanings: if a "friend" edge should matter more than a "follower" edge, attention weights learn that automatically. GraphSAGE is essential for production systems because it supports inductive inference — you can train on one graph and run inference on a completely unseen graph. Graph Transformers trade quadratic memory for the ability to capture long-range dependencies that shallow message passing cannot reach.

Training Pipeline with PyTorch Geometric

PyTorch Geometric (PyG) is the de facto standard for GNN training. Here is a complete training loop for node classification on the Cora citation network:

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

# 1. Load data
dataset = Planetoid(root="/tmp/cora", name="Cora")
data = dataset[0]

# 2. Define model
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)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training, p=0.5)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

model = GCN(
    in_channels=dataset.num_features,
    hidden_channels=16,
    out_channels=dataset.num_classes,
)
optimiser = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# 3. Training loop
def train():
    model.train()
    optimiser.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimiser.step()
    return float(loss)

# 4. Evaluation
@torch.no_grad()
def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        acc = (pred[mask] == data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs

for epoch in range(200):
    loss = train()
    train_acc, val_acc, test_acc = test()
    if epoch % 20 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss:.4f} | "
              f"Train: {train_acc:.3f} | Val: {val_acc:.3f} | Test: {test_acc:.3f}")

At 200 epochs on Cora (2,708 nodes, 5,429 edges), this GCN reaches roughly 81% test accuracy. The same architecture scaled to 100k nodes with GraphSAGE's neighbour sampling runs in minutes on a single GPU.

GNNs for Knowledge Graphs

GNNs have a natural overlap with knowledge graphs. A knowledge graph is a heterogeneous directed graph — nodes have types (Person, Company, Patent) and edges have labelled roles (WORKS_AT, FILES, CITES). Standard GNNs assume a single edge type with undirected semantics, so applying them to knowledge graphs requires extensions:

R-GCN (Schlichtkrull et al., 2018) maintains a separate transformation matrix per relation type. This works well for small-to-medium knowledge graphs but scales poorly — O(R × d²) parameters where R is the number of relation types and d is the hidden dimension.
CompGCN (Vashishth et al., 2020) composes edge embeddings with node embeddings during message passing, reducing the parameter count. It handles 10× more relation types than R-GCN for the same budget.
Link prediction in knowledge graphs uses a scoring function (DistMult, ComplEx, ConvE) on the node embeddings produced by the GNN. The model learns to score existing triples (h, r, t) higher than corrupted ones — the same loss structure used by TransE and its successors.

The result is a model that can predict missing links in your knowledge graph: "Which companies are likely to form partnerships?" or "What technologies might a given organisation develop next?"

Production Deployment Considerations

Deploying GNNs to production is harder than training them. Four challenges dominate:

Challenge	Problem	Mitigation
Inference on evolving graphs	Nodes and edges are added constantly; re-training is expensive	GraphSAGE inductive inference; streaming feature updates
Neighbourhood explosion	Full-batch adjacency on a 10M-node graph does not fit in GPU memory	Neighbour sampling (GraphSAINT, ClusterGCN); distributed mini-batch training
Feature staleness	Node features drift over time, degrading prediction quality	Online feature stores (Feast, Tecton) with TTL-based recomputation
Explainability	Regulators require reasons for edge-level predictions	GNNExplainer, Integrated Gradients on graph structure; subgraph-level explanations

Neighbour sampling is the single most important optimisation for production GNNs. Instead of materialising the full computation graph, each mini-batch samples a fixed-size neighbourhood per node — typically 10–25 neighbours at 2–3 hops. Tools like PyG's NeighborSampler and DGL's MultiLayerNeighborSampler handle this natively. For a graph of 10M nodes with average degree 50, sampling reduces per-batch memory from terabytes to megabytes.

For production inference, frameworks like TorchServe and NVIDIA Triton serve PyG models with GPU batching. A well-optimised GraphSAGE model with neighbour sampling can achieve sub-100ms p99 latency on graphs of 1M+ nodes with a single V100 GPU.

From Training to Inference: Exporting Your GNN

The training loop above produces a model in memory. For production, you need a serialised artefact that can be versioned and served independently of the training code. TorchScript traces the model's forward pass into a self-contained graph that runs without Python dependencies:

import torch

# Trace the trained GCN into TorchScript
model.eval()
traced = torch.jit.trace(model, (data.x, data.edge_index))
traced.save("gnn_cora.pt")

# Load and run inference — no PyG import needed
loaded = torch.jit.load("gnn_cora.pt")
loaded.eval()
with torch.no_grad():
    logits = loaded(data.x, data.edge_index)
    predictions = logits.argmax(dim=1)
    print(f"Predicted {predictions.sum():,} nodes")

For GPU serving at production scale, export to ONNX and serve via NVIDIA Triton or ONNX Runtime — both support dynamic batching and model versioning. TorchServe handles CPU deployments with integrated metrics and A/B testing. The neighbour sampling strategies discussed earlier combine naturally with these serving frameworks: GraphSAGE models exported to ONNX and served on Triton achieve sub-100ms p99 latency on graphs exceeding one million nodes.

Beyond Grids and Sequences

This article surveys the major GNN architectures, shows you how to train them with PyTorch Geometric, and covers what it takes to deploy them to production.

The Message-Passing Paradigm

h_v^(0) = x_v                          (initial node features)

h_v^(k) = UPDATE(
    h_v^(k-1),
    AGGREGATE({ h_u^(k-1) for every neighbour u of v })
)

GNN Architectures Compared

Architecture	Aggregation	Key Innovation	Best For	Limitations
GCN (Kipf & Welling, 2017)	Normalised mean	Symmetric normalisation of adjacency matrix	Citation networks, node classification	Treats all neighbours equally
GraphSAGE (Hamilton et al., 2017)	Mean / LSTM / Pool	Inductive — learns for unseen nodes; supports sampling	Large graphs, evolving graphs	Higher memory than GCN with LSTM aggregator
GAT (Veličković et al., 2018)	Attention-weighted sum	Self-attention over neighbours — learns which neighbours matter	Heterogeneous graphs, relationship-rich domains	O(E) attention computation
GIN (Xu et al., 2019)	Sum + MLP	Maximally expressive under WL-test	Graph classification, isomorphism detection	Deeper networks still over-smooth
Graph Transformer (Dwivedi & Bresson, 2021)	Full attention	Node positional encoding + transformer attention	Large graphs with long-range dependencies	O(V²) attention — does not scale beyond ~5k nodes

Training Pipeline with PyTorch Geometric

PyTorch Geometric (PyG) is the de facto standard for GNN training. Here is a complete training loop for node classification on the Cora citation network:

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

# 1. Load data
dataset = Planetoid(root="/tmp/cora", name="Cora")
data = dataset[0]

# 2. Define model
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)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training, p=0.5)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

model = GCN(
    in_channels=dataset.num_features,
    hidden_channels=16,
    out_channels=dataset.num_classes,
)
optimiser = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# 3. Training loop
def train():
    model.train()
    optimiser.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimiser.step()
    return float(loss)

# 4. Evaluation
@torch.no_grad()
def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        acc = (pred[mask] == data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs

for epoch in range(200):
    loss = train()
    train_acc, val_acc, test_acc = test()
    if epoch % 20 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss:.4f} | "
              f"Train: {train_acc:.3f} | Val: {val_acc:.3f} | Test: {test_acc:.3f}")

GNNs for Knowledge Graphs

R-GCN (Schlichtkrull et al., 2018) maintains a separate transformation matrix per relation type. This works well for small-to-medium knowledge graphs but scales poorly — O(R × d²) parameters where R is the number of relation types and d is the hidden dimension.
CompGCN (Vashishth et al., 2020) composes edge embeddings with node embeddings during message passing, reducing the parameter count. It handles 10× more relation types than R-GCN for the same budget.
Link prediction in knowledge graphs uses a scoring function (DistMult, ComplEx, ConvE) on the node embeddings produced by the GNN. The model learns to score existing triples (h, r, t) higher than corrupted ones — the same loss structure used by TransE and its successors.

The result is a model that can predict missing links in your knowledge graph: "Which companies are likely to form partnerships?" or "What technologies might a given organisation develop next?"

Production Deployment Considerations

Deploying GNNs to production is harder than training them. Four challenges dominate:

Challenge	Problem	Mitigation
Inference on evolving graphs	Nodes and edges are added constantly; re-training is expensive	GraphSAGE inductive inference; streaming feature updates
Neighbourhood explosion	Full-batch adjacency on a 10M-node graph does not fit in GPU memory	Neighbour sampling (GraphSAINT, ClusterGCN); distributed mini-batch training
Feature staleness	Node features drift over time, degrading prediction quality	Online feature stores (Feast, Tecton) with TTL-based recomputation
Explainability	Regulators require reasons for edge-level predictions	GNNExplainer, Integrated Gradients on graph structure; subgraph-level explanations

From Training to Inference: Exporting Your GNN

import torch

# Trace the trained GCN into TorchScript
model.eval()
traced = torch.jit.trace(model, (data.x, data.edge_index))
traced.save("gnn_cora.pt")

# Load and run inference — no PyG import needed
loaded = torch.jit.load("gnn_cora.pt")
loaded.eval()
with torch.no_grad():
    logits = loaded(data.x, data.edge_index)
    predictions = logits.argmax(dim=1)
    print(f"Predicted {predictions.sum():,} nodes")

Beyond Grids and Sequences

The Message-Passing Paradigm

GNN Architectures Compared

Training Pipeline with PyTorch Geometric

GNNs for Knowledge Graphs

Production Deployment Considerations

From Training to Inference: Exporting Your GNN

Further Reading

Graph Neural Networks: Architectures, Training, and Production Deployment

Beyond Grids and Sequences

The Message-Passing Paradigm

GNN Architectures Compared

Training Pipeline with PyTorch Geometric

GNNs for Knowledge Graphs

Production Deployment Considerations

From Training to Inference: Exporting Your GNN

Further Reading

Beyond Grids and Sequences

The Message-Passing Paradigm

GNN Architectures Compared

Training Pipeline with PyTorch Geometric

GNNs for Knowledge Graphs

Production Deployment Considerations

From Training to Inference: Exporting Your GNN

Further Reading

Footnotes

Never miss a deep-dive

Beyond Grids and Sequences

The Message-Passing Paradigm

GNN Architectures Compared

Training Pipeline with PyTorch Geometric

GNNs for Knowledge Graphs

Production Deployment Considerations

From Training to Inference: Exporting Your GNN

Further Reading

Footnotes

Never miss a deep-dive