Chapter 4: Real-World Applications

1. Molecular Property Prediction (Application to Drug Discovery)

In drug discovery, if the properties of candidate compounds (solubility, toxicity, membrane permeability, etc.) can be predicted before running experiments, development cost and time can be reduced significantly. Because molecules have a natural graph structure — atoms as nodes and chemical bonds as edges — GNNs are exceptionally well suited to molecular property prediction tasks.

Representing Molecules as Graphs

When treating a molecule as a graph, the typical correspondence is as follows.

In this section, we'll use the ESOL dataset from the MoleculeNet benchmark to implement a GNN model that predicts a molecule's aqueous solubility (log solubility) as a regression task. ESOL consists of 1,128 small molecules, and atom/bond features are generated automatically from SMILES notation (a string-based representation of molecular structure).

pip install rdkit

Loading the ESOL Dataset

import torch
from torch_geometric.datasets import MoleculeNet

dataset = MoleculeNet(root='/tmp/ESOL', name='ESOL')

print(f"Dataset: {dataset}")
print(f"Number of molecules: {len(dataset)}")
print(f"Number of node features: {dataset.num_node_features}")
print(f"Number of edge features: {dataset.num_edge_features}")

sample = dataset[0]
print(f"\nSample molecule: {sample}")
print(f"SMILES: {sample.smiles}")
print(f"Target (log solubility): {sample.y.item():.4f}")

Sample Output:

Dataset: ESOL(1128)
Number of molecules: 1128
Number of node features: 9
Number of edge features: 3

Sample molecule: Data(x=[32, 9], edge_index=[2, 68], edge_attr=[68, 3], y=[1, 1], smiles='OCC3OC(OCC2OC(OC(C#N)c1ccccc1)C(O)C(O)C2O)C(O)C(O)C3O')
SMILES: OCC3OC(OCC2OC(OC(C#N)c1ccccc1)C(O)C(O)C2O)C(O)C(O)C3O
Target (log solubility): -0.7700

A Molecular Property Prediction Model Using GIN

Molecular property prediction commonly uses a layer called GIN (Graph Isomorphism Network). GIN is theoretically designed to have the same discriminative power as the Weisfeiler-Lehman graph isomorphism test, making it well suited to capturing differences in small, dense structures like molecules.

$$\mathbf{x}_i^{(k+1)} = \text{MLP}^{(k)}\left((1 + \epsilon^{(k)}) \cdot \mathbf{x}_i^{(k)} + \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j^{(k)}\right)$$

import torch
import torch.nn.functional as F
from torch.nn import Linear, Sequential, ReLU, BatchNorm1d
from torch_geometric.nn import GINConv, global_mean_pool

class MolecularGIN(torch.nn.Module):
    def __init__(self, num_features, hidden_channels):
        super().__init__()
        nn1 = Sequential(Linear(num_features, hidden_channels), ReLU(),
                          Linear(hidden_channels, hidden_channels), BatchNorm1d(hidden_channels))
        self.conv1 = GINConv(nn1)

        nn2 = Sequential(Linear(hidden_channels, hidden_channels), ReLU(),
                          Linear(hidden_channels, hidden_channels), BatchNorm1d(hidden_channels))
        self.conv2 = GINConv(nn2)

        self.lin1 = Linear(hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)  # 1D output for the regression task

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()

        # Aggregate the whole graph into a single vector (graph-level pooling)
        x = global_mean_pool(x, batch)

        x = self.lin1(x).relu()
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.lin2(x)
        return x

Training and Evaluation

import torch
from torch_geometric.datasets import MoleculeNet
from torch_geometric.loader import DataLoader

dataset = MoleculeNet(root='/tmp/ESOL', name='ESOL')
dataset = dataset.shuffle()

# Convert features to float (some datasets store them as integers)
for data in dataset:
    data.x = data.x.float()

train_size = int(len(dataset) * 0.8)
train_dataset = dataset[:train_size]
test_dataset = dataset[train_size:]

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MolecularGIN(num_features=dataset.num_node_features, hidden_channels=64).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.MSELoss()

def train():
    model.train()
    total_loss = 0
    for batch in train_loader:
        batch = batch.to(device)
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index, batch.batch)
        loss = criterion(out.view(-1), batch.y.view(-1))
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * batch.num_graphs
    return total_loss / len(train_loader.dataset)

@torch.no_grad()
def evaluate(loader):
    model.eval()
    total_mae = 0
    for batch in loader:
        batch = batch.to(device)
        out = model(batch.x, batch.edge_index, batch.batch)
        total_mae += (out.view(-1) - batch.y.view(-1)).abs().sum().item()
    return total_mae / len(loader.dataset)

for epoch in range(1, 101):
    train_loss = train()
    if epoch % 20 == 0:
        test_mae = evaluate(test_loader)
        print(f'Epoch {epoch:03d}, Train MSE: {train_loss:.4f}, Test MAE: {test_mae:.4f}')

Sample Output:

Epoch 020, Train MSE: 1.1832, Test MAE: 0.8214
Epoch 040, Train MSE: 0.7145, Test MAE: 0.6839
Epoch 060, Train MSE: 0.5203, Test MAE: 0.6215
Epoch 080, Train MSE: 0.4109, Test MAE: 0.5981
Epoch 100, Train MSE: 0.3542, Test MAE: 0.5820

2. Citation Network Analysis

In Chapter 1, we performed node classification using the Cora dataset. In this section, we'll take citation network analysis a step further by implementing (1) qualitative analysis through node embedding visualization, and (2) link prediction, which predicts citation relationships that don't yet exist. Link prediction can be applied directly to recommending related papers, e.g., "this paper and that paper are likely to have a citation relationship in the future."

Designing the Link Prediction Task

In link prediction, we hide a portion of the known edges and train a model to predict the probability of an edge's existence from node embeddings. PyG's RandomLinkSplit can split edges into training, validation, and test sets, and automatically generates negative examples (non-existent edges).

import torch
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import RandomLinkSplit, NormalizeFeatures
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling

dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]

split = RandomLinkSplit(num_val=0.05, num_test=0.1, is_undirected=True,
                         add_negative_train_samples=False)
train_data, val_data, test_data = split(data)

print(f"Train edges: {train_data.edge_label_index.size(1)}")
print(f"Val edges (pos+neg): {val_data.edge_label_index.size(1)}")
print(f"Test edges (pos+neg): {test_data.edge_label_index.size(1)}")

Encoder and Link Prediction Model

class LinkPredictionGCN(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 encode(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

    def decode(self, z, edge_label_index):
        src, dst = edge_label_index
        return (z[src] * z[dst]).sum(dim=-1)  # Compute edge score via dot product

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LinkPredictionGCN(dataset.num_features, 64, 32).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

train_data = train_data.to(device)
val_data = val_data.to(device)
test_data = test_data.to(device)

def train_step():
    model.train()
    optimizer.zero_grad()
    z = model.encode(train_data.x, train_data.edge_index)

    # Sample the same number of negative edges (non-existent edges) as training edges
    neg_edge_index = negative_sampling(
        edge_index=train_data.edge_index,
        num_nodes=train_data.num_nodes,
        num_neg_samples=train_data.edge_label_index.size(1),
        method='sparse')

    edge_label_index = torch.cat([train_data.edge_label_index, neg_edge_index], dim=-1)
    edge_label = torch.cat([
        train_data.edge_label,
        train_data.edge_label.new_zeros(neg_edge_index.size(1))
    ], dim=0)

    out = model.decode(z, edge_label_index)
    loss = F.binary_cross_entropy_with_logits(out, edge_label)
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def eval_step(eval_data):
    model.eval()
    z = model.encode(eval_data.x, eval_data.edge_index)
    out = model.decode(z, eval_data.edge_label_index).sigmoid()
    return roc_auc_score(eval_data.edge_label.cpu().numpy(), out.cpu().numpy())

for epoch in range(1, 101):
    loss = train_step()
    if epoch % 20 == 0:
        val_auc = eval_step(val_data)
        print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, Val AUC: {val_auc:.4f}')

test_auc = eval_step(test_data)
print(f'\nTest AUC: {test_auc:.4f}')

Sample Output:

Epoch 020, Loss: 0.5842, Val AUC: 0.8215
Epoch 040, Loss: 0.4931, Val AUC: 0.8603
Epoch 060, Loss: 0.4520, Val AUC: 0.8791
Epoch 080, Loss: 0.4288, Val AUC: 0.8874
Epoch 100, Loss: 0.4109, Val AUC: 0.8902

Test AUC: 0.8856

An AUC (Area Under the Curve, the area under the ROC curve) of around 0.89 indicates that the model is doing a fairly good job of distinguishing positive examples (actual citations) from negative examples (non-existent citations). This mechanism can be applied directly to features such as "recommended papers for people reading this paper."

Visualizing Node Embeddings

Compressing the embedding space learned by the trained model down to two dimensions and visualizing it lets us qualitatively check whether papers from the same field cluster together.

import matplotlib
matplotlib.use('Agg')  # Backend for saving to file without on-screen display
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

model.eval()
with torch.no_grad():
    z = model.encode(data.x.to(device), data.edge_index.to(device)).cpu().numpy()

pca = PCA(n_components=2)
z_2d = pca.fit_transform(z)

plt.figure(figsize=(8, 6))
scatter = plt.scatter(z_2d[:, 0], z_2d[:, 1], c=data.y.numpy(), cmap='tab10', s=15)
plt.legend(*scatter.legend_elements(), title="Class", loc='best')
plt.title('PCA Visualization of Cora Node Embeddings')
plt.xlabel('Principal Component 1')
plt.ylabel('Principal Component 2')
plt.tight_layout()
plt.savefig('cora_embeddings.png', dpi=150)
print("Saved embedding visualization to cora_embeddings.png")

3. Recommendation Systems

Recommendation systems represent user-item interactions (purchases, views, ratings, etc.) as a bipartite graph, which is exactly the kind of structure GNNs excel at. By performing message passing between user nodes and item nodes, information about "items favored by users with similar tastes" propagates naturally.

Building the Bipartite Graph

Here, for training purposes, we'll create synthetic data that mimics user-item interactions. We place user nodes and item nodes in the same index space (item indices are offset by the number of users) and treat them as a single homogeneous graph.

import torch

torch.manual_seed(0)

num_users = 300
num_items = 150
num_raw_interactions = 3000

user_ids = torch.randint(0, num_users, (num_raw_interactions,))
item_ids = torch.randint(0, num_items, (num_raw_interactions,))

# Remove duplicate (user, item) pairs
interactions = torch.unique(torch.stack([user_ids, item_ids], dim=1), dim=0)
num_interactions = interactions.size(0)
print(f"Number of unique interactions: {num_interactions}")

item_node_offset = num_users
edge_user = interactions[:, 0]
edge_item = interactions[:, 1] + item_node_offset
num_nodes = num_users + num_items

# Split into train/test with an 8:2 ratio
perm = torch.randperm(num_interactions)
train_size = int(num_interactions * 0.8)
train_idx, test_idx = perm[:train_size], perm[train_size:]

train_user, train_item = edge_user[train_idx], edge_item[train_idx]
test_user, test_item = edge_user[test_idx], edge_item[test_idx]

# Build the message-passing graph only from training interactions (bidirectional)
train_edge_index = torch.cat([
    torch.stack([train_user, train_item], dim=0),
    torch.stack([train_item, train_user], dim=0),
], dim=1)

print(f"Number of training edges (including both directions): {train_edge_index.size(1)}")

Implementing the Graph-Based Recommender

We give each user and item a learnable initial embedding, aggregate neighborhood information with a GraphSAGE layer, and then compute a score via dot product. For training, we use the BPR loss (Bayesian Personalized Ranking Loss), optimizing so that the score of "an item the user actually interacted with" is higher than the score of "a randomly chosen item the user did not interact with."

$$\mathcal{L}_{\text{BPR}} = -\sum_{(u, i, j)} \log \sigma(\hat{y}_{ui} - \hat{y}_{uj})$$

import torch.nn.functional as F
from torch_geometric.nn import SAGEConv

class GraphRecommender(torch.nn.Module):
    def __init__(self, num_nodes, embedding_dim, hidden_dim):
        super().__init__()
        self.embedding = torch.nn.Embedding(num_nodes, embedding_dim)
        self.conv1 = SAGEConv(embedding_dim, hidden_dim)
        self.conv2 = SAGEConv(hidden_dim, hidden_dim)

    def forward(self, edge_index):
        x = self.embedding.weight
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

    def score(self, z, users, items):
        return (z[users] * z[items]).sum(dim=-1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphRecommender(num_nodes, embedding_dim=32, hidden_dim=32).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

train_edge_index_d = train_edge_index.to(device)
train_user_d = train_user.to(device)
train_item_d = train_item.to(device)

def sample_negative_items(users):
    # Sample random items as negatives, one per user
    return torch.randint(0, num_items, (users.size(0),), device=users.device) + item_node_offset

model.train()
for epoch in range(1, 51):
    optimizer.zero_grad()
    z = model(train_edge_index_d)

    pos_scores = model.score(z, train_user_d, train_item_d)
    neg_items = sample_negative_items(train_user_d)
    neg_scores = model.score(z, train_user_d, neg_items)

    loss = -F.logsigmoid(pos_scores - neg_scores).mean()
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f'Epoch {epoch:02d}, BPR Loss: {loss.item():.4f}')

Evaluation with Recall@K

@torch.no_grad()
def recall_at_k(z, test_user, test_item, train_user, train_item, k=10, num_eval_users=50):
    unique_test_users = torch.unique(test_user)[:num_eval_users]
    hits, total = 0, 0

    for u in unique_test_users:
        true_items = test_item[test_user == u]
        if true_items.numel() == 0:
            continue

        known_items = train_item[train_user == u] - item_node_offset
        item_embeddings = z[item_node_offset: item_node_offset + num_items]
        scores = (z[u] * item_embeddings).sum(dim=-1)
        scores[known_items] = -1e9  # Exclude known interactions from recommendation candidates

        topk_items = torch.topk(scores, k).indices + item_node_offset
        hit = any(t.item() in topk_items.tolist() for t in true_items)
        hits += int(hit)
        total += 1

    return hits / total if total > 0 else 0.0

model.eval()
z = model(train_edge_index_d)
recall = recall_at_k(z, test_user.to(device), test_item.to(device),
                      train_user_d, train_item_d, k=10)
print(f'Recall@10: {recall:.4f}')

4. Efficient Processing of Large-Scale Graphs

For a graph with a few thousand nodes, like Cora, full-batch training — loading all nodes onto the GPU at once — is feasible. In practice, however, it's not uncommon for graphs to have millions to hundreds of millions of nodes. For such large-scale graphs, it becomes physically impossible to load all nodes and their features into memory.

Mini-Batch Training and Neighbor Sampling

The solution to this problem is mini-batch training based on neighbor sampling. In each mini-batch, we select target nodes (seed nodes) and sample only a fixed number of neighboring nodes per layer, using only the necessary subgraph for computation. In PyG, NeighborLoader implements this.

To get a feel for large-scale graphs, we'll generate a synthetic graph with 5,000 nodes and implement mini-batch training with NeighborLoader. The same code pattern applies directly to real large-scale graphs with millions of nodes (e.g., ogbn-products or Reddit).

import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.utils import erdos_renyi_graph
from torch_geometric.loader import NeighborLoader
from torch_geometric.nn import SAGEConv

torch.manual_seed(42)

num_nodes = 5000
num_features = 32
num_classes = 5

edge_index = erdos_renyi_graph(num_nodes, edge_prob=0.002, directed=False)
x = torch.randn(num_nodes, num_features)
y = torch.randint(0, num_classes, (num_nodes,))

perm = torch.randperm(num_nodes)
train_mask = torch.zeros(num_nodes, dtype=torch.bool)
val_mask = torch.zeros(num_nodes, dtype=torch.bool)
test_mask = torch.zeros(num_nodes, dtype=torch.bool)
train_mask[perm[:3000]] = True
val_mask[perm[3000:4000]] = True
test_mask[perm[4000:]] = True

data = Data(x=x, edge_index=edge_index, y=y,
            train_mask=train_mask, val_mask=val_mask, test_mask=test_mask)

print(data)
print(f"Average degree: {data.num_edges / data.num_nodes:.2f}")

Mini-Batch Training with NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[10, 10],  # Number of neighbors to sample per layer (10 for layer 1, 10 for layer 2)
    batch_size=256,
    input_nodes=data.train_mask,
    shuffle=True,
)

class SAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, out_channels)

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SAGE(num_features, hidden_channels=64, out_channels=num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

model.train()
for epoch in range(1, 6):
    total_loss = 0
    for batch in train_loader:
        batch = batch.to(device)
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index)

        # In NeighborLoader batches, the first batch.batch_size nodes are the seed nodes
        loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * batch.batch_size

    avg_loss = total_loss / int(data.train_mask.sum())
    print(f'Epoch {epoch:02d}, Loss: {avg_loss:.4f}')

Sample Output:

Epoch 01, Loss: 1.6132
Epoch 02, Loss: 1.5890
Epoch 03, Loss: 1.5701
Epoch 04, Loss: 1.5523
Epoch 05, Loss: 1.5388

5. Model Interpretability and Visualization

While GNNs can achieve high predictive accuracy, they tend to be black-box models that are hard to explain in terms of "why did it make this prediction." In domains involving decision-making, such as drug discovery and healthcare, model interpretability is a practical necessity. PyG provides GNNExplainer, a leading explanation method, through the torch_geometric.explain module.

Visualizing Prediction Rationale with GNNExplainer

GNNExplainer reveals which edges and which features were important for the prediction of a specific node by learning small masks (weights).

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
from torch_geometric.explain import Explainer, GNNExplainer

dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

class GCN(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = GCNConv(num_features, 16)
        self.conv2 = GCNConv(16, num_classes)

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN(dataset.num_features, dataset.num_classes).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# Quick training (shorter epoch count for the interpretability demo)
model.train()
for epoch in range(100):
    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()

print(f"Final train loss: {loss.item():.4f}")

explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config=dict(
        mode='multiclass_classification',
        task_level='node',
        return_type='raw',
    ),
)

# Explain the prediction rationale for node 10
node_index = 10
explanation = explainer(data.x, data.edge_index, index=node_index)

print(f"Explained node: {node_index}")
print(f"Important feature mask shape: {explanation.node_mask.shape}")
print(f"Important edge mask shape: {explanation.edge_mask.shape}")

# Top 5 most important features
top_features = explanation.node_mask[node_index].topk(5)
print(f"\nTop important feature indices: {top_features.indices.tolist()}")
print(f"Importance scores: {[round(v, 4) for v in top_features.values.tolist()]}")

# Top 5 most important edges
top_edges = explanation.edge_mask.topk(5)
print(f"\nTop important edge indices: {top_edges.indices.tolist()}")

Sample Output:

Explained node: 10
Important feature mask shape: torch.Size([2708, 1433])
Important edge mask shape: torch.Size([10556])

Top important feature indices: [19, 142, 501, 88, 973]
Importance scores: [0.8123, 0.7654, 0.7211, 0.6987, 0.6544]

Top important edge indices: [3421, 128, 5502, 891, 2077]

Visualizing the Explanation Subgraph

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import networkx as nx
from torch_geometric.utils import to_networkx

# Visualize a subgraph extracted from the top-important edges
threshold = explanation.edge_mask.topk(15).values.min()
important_edge_mask = explanation.edge_mask >= threshold
important_edges = data.edge_index[:, important_edge_mask].cpu()

involved_nodes = torch.unique(important_edges).tolist()
if node_index not in involved_nodes:
    involved_nodes.append(node_index)

subgraph = nx.Graph()
subgraph.add_nodes_from(involved_nodes)
for src, dst in important_edges.t().tolist():
    subgraph.add_edge(src, dst)

colors = ['#e74c3c' if n == node_index else '#3498db' for n in subgraph.nodes()]

plt.figure(figsize=(7, 6))
nx.draw(subgraph, node_color=colors, with_labels=True, node_size=400,
        font_size=8, font_color='white')
plt.title(f'Subgraph Important for the Prediction of Node {node_index}')
plt.savefig('gnn_explanation.png', dpi=150, bbox_inches='tight')
print("Saved explanation subgraph to gnn_explanation.png")

6. Deploying to Production

Even after a highly accurate GNN model is complete at the research/experimentation stage, running it reliably in a production environment involves many considerations beyond those at training time. Here, using the molecular property prediction model from Section 1 (the ESOL model) as an example, we'll walk through practical deployment patterns.

Saving and Loading the Model for Inference

import torch

# Save only the model weights (recommended approach)
torch.save(model.state_dict(), 'molecular_gin_esol.pt')

# When loading, redefine the model architecture before loading the weights
loaded_model = MolecularGIN(num_features=9, hidden_channels=64)
loaded_model.load_state_dict(torch.load('molecular_gin_esol.pt', map_location='cpu'))
loaded_model.eval()

print("Model loading complete")

Inference Pipeline from SMILES to Prediction

In production, the same preprocessing used at training time (converting SMILES strings into molecular graphs) must be reproduced at inference time as well. Bundling the pre- and post-processing into a single function helps prevent mistakes during deployment.

import torch
from torch_geometric.datasets import MoleculeNet
from torch_geometric.loader import DataLoader

# Reuses MoleculeNet's feature extraction logic to apply
# the exact same preprocessing used during training to a single molecule
def predict_solubility(model, molecule_data, device='cpu'):
    """Helper function that performs inference for a single molecule.

    Args:
        model: A trained MolecularGIN model (already in eval mode)
        molecule_data: torch_geometric.data.Data (containing x, edge_index)
        device: The device to use for inference

    Returns:
        float: The predicted log solubility
    """
    model = model.to(device)
    molecule_data = molecule_data.to(device)
    with torch.no_grad():
        x = molecule_data.x.float()
        # For a single graph, we need a batch vector indicating all nodes belong to the same graph ID
        batch = torch.zeros(x.size(0), dtype=torch.long, device=device)
        pred = model(x, molecule_data.edge_index, batch)

    return pred.item()

# Verify the inference pipeline using a molecule from the validation dataset
dataset = MoleculeNet(root='/tmp/ESOL', name='ESOL')
sample_molecule = dataset[0]

prediction = predict_solubility(loaded_model, sample_molecule)
print(f"SMILES: {sample_molecule.smiles}")
print(f"Predicted log solubility: {prediction:.4f}")
print(f"Actual value: {sample_molecule.y.item():.4f}")

Production Deployment Checklist

Exercises

import torch
from torch.nn import Linear, Sequential, ReLU
from torch_geometric.nn import GINEConv, global_mean_pool

class MolecularGINE(torch.nn.Module):
    def __init__(self, num_features, num_edge_features, hidden_channels):
        super().__init__()
        nn1 = Sequential(Linear(hidden_channels, hidden_channels), ReLU(),
                          Linear(hidden_channels, hidden_channels))
        # GINEConv requires the input feature and edge feature dimensions to match,
        # so we align them beforehand with a linear layer
        self.node_proj = Linear(num_features, hidden_channels)
        self.edge_proj = Linear(num_edge_features, hidden_channels)
        self.conv1 = GINEConv(nn1)

        self.lin1 = Linear(hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)

    def forward(self, x, edge_index, edge_attr, batch):
        x = self.node_proj(x)
        edge_attr = self.edge_proj(edge_attr)
        x = self.conv1(x, edge_index, edge_attr).relu()
        x = global_mean_pool(x, batch)
        x = self.lin1(x).relu()
        x = self.lin2(x)
        return x

# Usage example:
# model = MolecularGINE(num_features=9, num_edge_features=3, hidden_channels=64)
# out = model(batch.x.float(), batch.edge_index, batch.edge_attr.float(), batch.batch)

import time
import torch.nn.functional as F
from torch_geometric.loader import NeighborLoader

def run_experiment(data, num_neighbors, num_features, num_classes, device):
    loader = NeighborLoader(
        data, num_neighbors=num_neighbors, batch_size=256,
        input_nodes=data.train_mask, shuffle=True,
    )
    model = SAGE(num_features, 64, num_classes).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    start = time.time()
    model.train()
    total_loss = 0
    for batch in loader:
        batch = batch.to(device)
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index)
        loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * batch.batch_size
    elapsed = time.time() - start
    avg_loss = total_loss / int(data.train_mask.sum())
    return elapsed, avg_loss

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for neighbors in ([5, 5], [10, 10], [25, 25]):
    elapsed, avg_loss = run_experiment(data, neighbors, num_features, num_classes, device)
    print(f"num_neighbors={neighbors}: time={elapsed:.2f}s, loss={avg_loss:.4f}")

import torch

model.eval()
with torch.no_grad():
    pred = model(data.x, data.edge_index).argmax(dim=1)

correct_mask = (pred == data.y) & data.test_mask
incorrect_mask = (pred != data.y) & data.test_mask

correct_node = correct_mask.nonzero(as_tuple=True)[0][0].item()
incorrect_candidates = incorrect_mask.nonzero(as_tuple=True)[0]

if incorrect_candidates.numel() > 0:
    incorrect_node = incorrect_candidates[0].item()

    for label, node_idx in [("Correctly classified", correct_node), ("Misclassified", incorrect_node)]:
        explanation = explainer(data.x, data.edge_index, index=node_idx)
        top_edges = explanation.edge_mask.topk(5)
        print(f"\n[{label}] Node {node_idx}")
        print(f"  True label: {data.y[node_idx].item()}, Predicted: {pred[node_idx].item()}")
        print(f"  Important edge indices: {top_edges.indices.tolist()}")
else:
    print("No misclassified nodes found in the test set (this can happen when model accuracy is very high)")

Summary

In this chapter, we learned practical techniques for applying GNNs to real-world problems.

• ✅ Molecular property prediction (regression task) using MoleculeNet (ESOL) and GINConv
• ✅ Link prediction for citation networks using RandomLinkSplit and a GCN encoder, plus embedding visualization with PCA
• ✅ Graph-based recommendation using a bipartite graph, GraphSAGE, and BPR loss
• ✅ Mini-batch training for large-scale graphs using neighbor sampling with NeighborLoader
• ✅ Visualizing prediction rationale and improving model interpretability with GNNExplainer
• ✅ Building a model save/inference pipeline and practical considerations for production deployment

Reference Resources

• PyTorch Geometric Official Documentation
• PyG Explainability Official Tutorial
• MoleculeNet: Benchmark Datasets
• GIN Paper: How Powerful are Graph Neural Networks?
• GraphSAGE Paper: Inductive Representation Learning on Large Graphs