Chapter 3: Hands-on PyTorch Geometric - Implementing Molecular and Material Property Prediction

Experience the fastest route from data preparation → model training → evaluation with PyG. We also provide a checklist for when training is unstable.

💡 Note: First verify operation with small graphs and few epochs → then scale up by increasing the batch size and number of layers, in that order. Start with a small learning rate.

Learning Objectives

By reading this chapter, you will acquire the following: - Set up a PyTorch Geometric environment and master the use of GNN libraries - Implement a molecular property prediction model on the QM9 dataset - Perform crystal property prediction on Materials Project data - Apply best practices for model training - Visualize prediction results and evaluate performance

Reading time: 25-30 minutes Code examples: 10 Exercises: 3

3.1 Environment Setup: Installing PyTorch Geometric

3.1.1 What Is PyTorch Geometric?

PyTorch Geometric (PyG) is a library dedicated to graph neural networks that runs on top of PyTorch.

Key features: - 🚀 Fast: Efficient graph processing via GPU - 📦 Rich set of models: More than 30 models including GCN, GAT, GraphSAGE, and SchNet - 🧪 Datasets: QM9, ZINC, and OGB (Open Graph Benchmark) are built in - 🛠️ Flexibility: Custom layers and models can be implemented easily

3.1.2 Installation Steps

Option 1: Conda environment (recommended)

# 1. Create an environment with Python 3.9 or later
conda create -n gnn-env python=3.10
conda activate gnn-env

# 2. Install PyTorch (CUDA version recommended)
# For the CPU version:
conda install pytorch torchvision torchaudio cpuonly -c pytorch

# For the GPU version (CUDA 11.8):
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# 3. Install PyTorch Geometric
conda install pyg -c pyg

# 4. Additional libraries
pip install rdkit matplotlib seaborn pandas scikit-learn

Option 2: Installation via pip

# 1. Create a virtual environment
python -m venv gnn-env
source gnn-env/bin/activate  # macOS/Linux
# gnn-env\Scripts\activate  # Windows

# 2. Install PyTorch
pip install torch torchvision torchaudio

# 3. Install PyTorch Geometric
pip install torch-geometric

# 4. Dependency libraries
pip install torch-scatter torch-sparse torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cpu.html

# 5. Additional libraries
pip install rdkit matplotlib seaborn pandas scikit-learn

Option 3: Google Colab (no installation required)

# In Google Colab, run the following
!pip install torch-geometric
!pip install rdkit

3.1.3 Verifying the Installation

import torch
import torch_geometric
from torch_geometric.data import Data
from rdkit import Chem
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

print("===== Installation Check =====")
print(f"PyTorch version: {torch.__version__}")
print(f"PyTorch Geometric version: {torch_geometric.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# Create a simple graph to test
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
data = Data(x=x, edge_index=edge_index)

print(f"\nTest graph created successfully!")
print(f"Number of nodes: {data.num_nodes}")
print(f"Number of edges: {data.num_edges}")
print("✅ PyTorch Geometric environment setup complete!")

Expected output:

===== Installation Check =====
PyTorch version: 2.0.0
PyTorch Geometric version: 2.3.0
CUDA available: True
CUDA version: 11.8
GPU: NVIDIA GeForce RTX 3090

Test graph created successfully!
Number of nodes: 3
Number of edges: 4
✅ PyTorch Geometric environment setup complete!

3.1.4 Troubleshooting

3.2 PyTorch Geometric Basics: Data Structures and DataLoader

3.2.1 Structure of the Data Object

In PyTorch Geometric, a graph is represented by a Data object.

from torch_geometric.data import Data
import torch

# Represent an ethanol molecule (C2H5OH) as a graph
# C: carbon (nodes 0, 1)
# O: oxygen (node 2)
# H: hydrogen (nodes 3-7)

# Node features (using atomic numbers)
x = torch.tensor([
    [6],   # C (carbon)
    [6],   # C (carbon)
    [8],   # O (oxygen)
    [1],   # H (hydrogen)
    [1],   # H (hydrogen)
    [1],   # H (hydrogen)
    [1],   # H (hydrogen)
    [1],   # H (hydrogen)
], dtype=torch.float)

# Edge index (bonding relationships)
# Each bond is bidirectional (undirected graph)
edge_index = torch.tensor([
    [0, 1, 1, 0, 0, 2, 2, 0, 0, 3, 3, 0, 1, 4, 4, 1, 1, 5, 5, 1, 2, 6, 6, 2],
    [1, 0, 2, 2, 3, 0, 0, 3, 4, 1, 1, 4, 5, 1, 1, 5, 6, 2, 2, 6, 7, 2, 2, 7]
], dtype=torch.long)

# Edge features (bond type: 1 = single bond)
edge_attr = torch.ones(edge_index.size(1), 1)

# Molecule-level feature (target variable)
y = torch.tensor([[156.0]], dtype=torch.float)  # boiling point (°C)

# Create the Data object
ethanol = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)

print("===== Graph Representation of the Ethanol Molecule =====")
print(f"Number of nodes (atoms): {ethanol.num_nodes}")
print(f"Number of edges (bonds x 2): {ethanol.num_edges}")
print(f"Shape of node features: {ethanol.x.shape}")
print(f"Shape of edge index: {ethanol.edge_index.shape}")
print(f"Target variable (boiling point): {ethanol.y.item()} °C")

# Basic graph statistics
print(f"\n===== Graph Statistics =====")
print(f"Average degree (number of bonds): {ethanol.num_edges / ethanol.num_nodes:.2f}")
print(f"Isolated nodes: {ethanol.contains_isolated_nodes()}")
print(f"Self-loops: {ethanol.contains_self_loops()}")

Output:

===== Graph Representation of the Ethanol Molecule =====
Number of nodes (atoms): 8
Number of edges (bonds x 2): 24
Shape of node features: torch.Size([8, 1])
Shape of edge index: torch.Size([2, 24])
Target variable (boiling point): 156.0 °C

===== Graph Statistics =====
Average degree (number of bonds): 3.00
Isolated nodes: False
Self-loops: False

3.2.2 Converting from RDKit to a Graph

RDKit can create a molecule object from a SMILES string (a text representation of a molecule).

from rdkit import Chem
from rdkit.Chem import Draw
from torch_geometric.data import Data
import torch

def mol_to_graph(smiles):
    """
    Create a PyTorch Geometric Data object from a SMILES string

    Parameters:
    -----------
    smiles : str
        SMILES representation of the molecule

    Returns:
    --------
    data : torch_geometric.data.Data
        Graph data
    """
    # Create a molecule object from SMILES
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None

    # Node features (atomic properties)
    atom_features = []
    for atom in mol.GetAtoms():
        # One-hot encode the atomic number (C, N, O, F, other)
        atom_type = [0] * 5
        if atom.GetAtomicNum() == 6:    # C
            atom_type[0] = 1
        elif atom.GetAtomicNum() == 7:  # N
            atom_type[1] = 1
        elif atom.GetAtomicNum() == 8:  # O
            atom_type[2] = 1
        elif atom.GetAtomicNum() == 9:  # F
            atom_type[3] = 1
        else:
            atom_type[4] = 1

        # Add formal charge and aromaticity
        formal_charge = atom.GetFormalCharge()
        is_aromatic = int(atom.GetIsAromatic())

        atom_features.append(atom_type + [formal_charge, is_aromatic])

    x = torch.tensor(atom_features, dtype=torch.float)

    # Edge index (bonding relationships)
    edge_indices = []
    for bond in mol.GetBonds():
        i = bond.GetBeginAtomIdx()
        j = bond.GetEndAtomIdx()
        edge_indices += [[i, j], [j, i]]  # bidirectional because undirected graph

    edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()

    data = Data(x=x, edge_index=edge_index)
    return data, mol

# Test: convert several molecules to graphs
smiles_list = [
    ("C", "Methane"),
    ("CCO", "Ethanol"),
    ("c1ccccc1", "Benzene"),
    ("CC(=O)O", "Acetic acid"),
]

print("===== Converting SMILES to Graphs =====")
for smiles, name in smiles_list:
    data, mol = mol_to_graph(smiles)
    print(f"\n{name} ({smiles}):")
    print(f"  Number of nodes: {data.num_nodes}")
    print(f"  Number of edges: {data.num_edges}")
    print(f"  Node feature dimension: {data.x.shape[1]}")

# Visualize molecular structures
import matplotlib.pyplot as plt
from rdkit.Chem import Draw

fig, axes = plt.subplots(1, 4, figsize=(16, 4))
for i, (smiles, name) in enumerate(smiles_list):
    _, mol = mol_to_graph(smiles)
    img = Draw.MolToImage(mol, size=(300, 300))
    axes[i].imshow(img)
    axes[i].set_title(f"{name}\n{smiles}", fontsize=12)
    axes[i].axis('off')

plt.tight_layout()
plt.show()

3.2.3 Using the DataLoader

To process multiple graphs in batches, use the DataLoader.

from torch_geometric.data import Data, DataLoader
import torch

# Create a sample dataset (10 molecules)
dataset = []
for i in range(10):
    num_nodes = torch.randint(5, 15, (1,)).item()  # 5-14 atoms
    x = torch.randn(num_nodes, 7)  # node features (7 dimensions)

    # Generate random edges
    edge_index = torch.randint(0, num_nodes, (2, num_nodes * 2))

    # Target variable (e.g., HOMO-LUMO gap)
    y = torch.randn(1)

    data = Data(x=x, edge_index=edge_index, y=y)
    dataset.append(data)

# Create the DataLoader (batch size = 4)
loader = DataLoader(dataset, batch_size=4, shuffle=True)

print("===== Using the DataLoader =====")
print(f"Dataset size: {len(dataset)}")
print(f"Number of batches: {len(loader)}")

# Inspect the first batch
for batch in loader:
    print(f"\nFirst batch:")
    print(f"  Number of molecules in the batch: {batch.num_graphs}")
    print(f"  Total number of nodes: {batch.num_nodes}")
    print(f"  Total number of edges: {batch.num_edges}")
    print(f"  Shape of node features: {batch.x.shape}")
    print(f"  Batch index: {batch.batch}")
    print(f"  Shape of target variable: {batch.y.shape}")
    break

Example output:

===== Using the DataLoader =====
Dataset size: 10
Number of batches: 3

First batch:
  Number of molecules in the batch: 4
  Total number of nodes: 38
  Total number of edges: 76
  Shape of node features: torch.Size([38, 7])
  Batch index: tensor([0, 0, 0, ..., 3, 3, 3])
  Shape of target variable: torch.Size([4, 1])

Important: The batch tensor indicates which molecule each node belongs to (0, 0, 0, 1, 1, 2, 2, 2, 3, ...).

3.3 Molecular Property Prediction with the QM9 Dataset

3.3.1 Overview of the QM9 Dataset

QM9 is a quantum-chemistry calculation dataset of 134,000 small organic molecules.

Included properties: - HOMO (highest occupied molecular orbital energy) - LUMO (lowest unoccupied molecular orbital energy) - Band gap (HOMO-LUMO gap) - Dipole moment - Internal energy - Enthalpy, free energy, heat capacity, and others

3.3.2 Loading the QM9 Dataset

from torch_geometric.datasets import QM9
import torch

# Download the dataset (first time only, about 1 GB)
dataset = QM9(root='./data/QM9')

print("===== QM9 Dataset =====")
print(f"Number of molecules: {len(dataset)}")
print(f"Node feature dimension: {dataset.num_node_features}")
print(f"Edge feature dimension: {dataset.num_edge_features}")
print(f"Number of target variables: {dataset.num_classes}")

# Inspect the first molecule
data = dataset[0]
print(f"\nFirst molecule:")
print(f"  Number of atoms: {data.num_nodes}")
print(f"  Number of bonds: {data.num_edges // 2}")
print(f"  Node features: {data.x.shape}")
print(f"  Edge features: {data.edge_attr.shape}")
print(f"  Target variables (19 types): {data.y.shape}")

# Display some of the target variables
target_names = ['mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve',
                'U0', 'U', 'H', 'G', 'Cv']
print(f"\nKey property values:")
for i, name in enumerate(target_names):
    print(f"  {name}: {data.y[0, i].item():.4f}")

3.3.3 Implementing a Graph Convolutional Network (GCN)

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool

class GCN_QM9(torch.nn.Module):
    """
    Graph Convolutional Network for QM9 molecular property prediction

    Architecture:
    - 3 GCNConv layers
    - Global mean pooling
    - 2 fully connected layers
    """
    def __init__(self, num_node_features, num_classes, hidden_channels=64):
        super(GCN_QM9, self).__init__()

        # GCN layers
        self.conv1 = GCNConv(num_node_features, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)

        # Fully connected layers
        self.lin1 = torch.nn.Linear(hidden_channels, hidden_channels // 2)
        self.lin2 = torch.nn.Linear(hidden_channels // 2, num_classes)

    def forward(self, x, edge_index, batch):
        """
        Parameters:
        -----------
        x : torch.Tensor (num_nodes, num_node_features)
            Node features
        edge_index : torch.Tensor (2, num_edges)
            Edge index
        batch : torch.Tensor (num_nodes,)
            Batch index

        Returns:
        --------
        out : torch.Tensor (batch_size, num_classes)
            Predicted values
        """
        # GCN layer 1 (convolution + activation + dropout)
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.2, training=self.training)

        # GCN layer 2
        x = self.conv2(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.2, training=self.training)

        # GCN layer 3
        x = self.conv3(x, edge_index)
        x = F.relu(x)

        # Global pooling (aggregate node features to the molecule level)
        x = global_mean_pool(x, batch)

        # Fully connected layers
        x = self.lin1(x)
        x = F.relu(x)
        x = F.dropout(x, p=0.3, training=self.training)

        x = self.lin2(x)
        return x

# Instantiate the model
model = GCN_QM9(
    num_node_features=dataset.num_node_features,
    num_classes=1,  # predict only the HOMO-LUMO gap
    hidden_channels=64
)

print("===== GCN Model =====")
print(model)
print(f"\nNumber of parameters: {sum(p.numel() for p in model.parameters()):,}")

3.3.4 Training the Model

from torch_geometric.loader import DataLoader
from sklearn.model_selection import train_test_split
import time

# Shrink the dataset (for speed; in practice use the full data)
dataset = dataset[:10000]

# Set only the HOMO-LUMO gap (index=4) as the target variable
for data in dataset:
    data.y = data.y[:, 4:5]  # shape: (1, 1)

# Data split (80% train, 10% val, 10% test)
train_dataset = dataset[:8000]
val_dataset = dataset[8000:9000]
test_dataset = dataset[9000:]

# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# Device setup (use GPU if available, otherwise CPU)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Loss function and optimization algorithm
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)

# Training function
def train(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0

    for data in loader:
        data = data.to(device)
        optimizer.zero_grad()

        # Forward pass
        out = model(data.x, data.edge_index, data.batch)
        loss = criterion(out, data.y)

        # Backward pass
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * data.num_graphs

    return total_loss / len(loader.dataset)

# Validation function
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0

    with torch.no_grad():
        for data in loader:
            data = data.to(device)
            out = model(data.x, data.edge_index, data.batch)
            loss = criterion(out, data.y)
            total_loss += loss.item() * data.num_graphs

    return total_loss / len(loader.dataset)

# Training loop
epochs = 50
train_losses = []
val_losses = []
best_val_loss = float('inf')

print("===== Training Start =====")
start_time = time.time()

for epoch in range(1, epochs + 1):
    train_loss = train(model, train_loader, optimizer, criterion, device)
    val_loss = evaluate(model, val_loader, criterion, device)

    train_losses.append(train_loss)
    val_losses.append(val_loss)

    # Save the best model
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best_model_qm9.pt')

    if epoch % 10 == 0:
        print(f"Epoch {epoch:03d}, "
              f"Train Loss: {train_loss:.4f}, "
              f"Val Loss: {val_loss:.4f}")

training_time = time.time() - start_time
print(f"\nTraining complete! Elapsed time: {training_time:.2f} seconds")

# Load the best model
model.load_state_dict(torch.load('best_model_qm9.pt'))

# Evaluate on the test data
test_loss = evaluate(model, test_loader, criterion, device)
test_mae = test_loss ** 0.5  # use RMSE as an approximation of MAE

print(f"\n===== Test Performance =====")
print(f"Test Loss (MSE): {test_loss:.4f}")
print(f"Test MAE (approx): {test_mae:.4f} eV")

3.3.5 Visualizing the Learning Curve

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(10, 6))
ax.plot(train_losses, label='Train Loss', linewidth=2)
ax.plot(val_losses, label='Validation Loss', linewidth=2)
ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Loss (MSE)', fontsize=12)
ax.set_title('GCN Learning Curve (QM9 HOMO-LUMO Gap Prediction)', fontsize=14)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# Plot predicted vs. actual
model.eval()
all_preds = []
all_targets = []

with torch.no_grad():
    for data in test_loader:
        data = data.to(device)
        out = model(data.x, data.edge_index, data.batch)
        all_preds.append(out.cpu().numpy())
        all_targets.append(data.y.cpu().numpy())

all_preds = np.concatenate(all_preds)
all_targets = np.concatenate(all_targets)

fig, ax = plt.subplots(figsize=(8, 8))
ax.scatter(all_targets, all_preds, alpha=0.6, s=10)
ax.plot([all_targets.min(), all_targets.max()],
        [all_targets.min(), all_targets.max()],
        'r--', lw=2, label='Perfect prediction')
ax.set_xlabel('Actual value (eV)', fontsize=12)
ax.set_ylabel('Predicted value (eV)', fontsize=12)
ax.set_title('HOMO-LUMO Gap Prediction Results', fontsize=14)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)

# Compute the R² score
from sklearn.metrics import r2_score
r2 = r2_score(all_targets, all_preds)
mae = np.mean(np.abs(all_targets - all_preds))

ax.text(0.05, 0.95, f'R² = {r2:.3f}\nMAE = {mae:.3f} eV',
        transform=ax.transAxes, fontsize=12, verticalalignment='top',
        bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.tight_layout()
plt.show()

print(f"===== Final Performance =====")
print(f"R² score: {r2:.4f}")
print(f"MAE: {mae:.4f} eV")

3.4 Crystal Property Prediction with Materials Project Data

3.4.1 Graph Representation of Crystal Structures

Because crystals have periodic structures, they require different handling from molecules.

from pymatgen.core import Structure
from pymatgen.ext.matproj import MPRester
import torch
from torch_geometric.data import Data

def structure_to_graph(structure, cutoff=5.0):
    """
    Convert a pymatgen Structure object to a graph

    Parameters:
    -----------
    structure : pymatgen.core.Structure
        Crystal structure
    cutoff : float
        Distance cutoff for creating edges (Å)

    Returns:
    --------
    data : torch_geometric.data.Data
        Graph data
    """
    # Node features (atomic numbers)
    atomic_numbers = [site.specie.Z for site in structure]
    x = torch.tensor(atomic_numbers, dtype=torch.float).view(-1, 1)

    # Edge index and edge features (interatomic distances)
    edge_indices = []
    edge_attrs = []

    for i, site_i in enumerate(structure):
        for j, site_j in enumerate(structure):
            if i != j:
                distance = structure.get_distance(i, j)
                if distance < cutoff:
                    edge_indices.append([i, j])
                    edge_attrs.append([distance])

    edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
    edge_attr = torch.tensor(edge_attrs, dtype=torch.float)

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
    return data

# Retrieve Li compounds from Materials Project (sample)
# Note: an API key is required in practice
# API_KEY = "your_api_key_here"
# with MPRester(API_KEY) as mpr:
#     entries = mpr.query(
#         criteria={"elements": {"$all": ["Li"]}, "nelements": 2},
#         properties=["structure", "band_gap"]
#     )

# Sample data (LiCl crystal)
from pymatgen.core import Lattice, Structure

# LiCl rock-salt structure
lattice = Lattice.cubic(5.14)  # lattice constant
species = ["Li", "Li", "Li", "Li", "Cl", "Cl", "Cl", "Cl"]
coords = [
    [0, 0, 0], [0.5, 0.5, 0], [0.5, 0, 0.5], [0, 0.5, 0.5],
    [0.5, 0, 0], [0, 0.5, 0], [0, 0, 0.5], [0.5, 0.5, 0.5]
]
structure = Structure(lattice, species, coords)

# Convert to a graph
data = structure_to_graph(structure, cutoff=4.0)

print("===== Graph Representation of the LiCl Crystal =====")
print(f"Number of nodes (atoms): {data.num_nodes}")
print(f"Number of edges (atom pairs with distance < 4.0 Å): {data.num_edges}")
print(f"Node features: {data.x}")
print(f"\nStatistics of edge features (distances):")
print(f"  Minimum distance: {data.edge_attr.min().item():.2f} Å")
print(f"  Maximum distance: {data.edge_attr.max().item():.2f} Å")
print(f"  Average distance: {data.edge_attr.mean().item():.2f} Å")

3.4.2 Crystal Property Prediction Model (Crystal Graph Convolutional Network)

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_add_pool

class CGCN(torch.nn.Module):
    """
    Crystal Graph Convolutional Network
    Predicts the band gap of crystals
    """
    def __init__(self, num_node_features=1, hidden_channels=64):
        super(CGCN, self).__init__()

        # Node embedding layer
        self.embedding = torch.nn.Linear(num_node_features, hidden_channels)

        # GCN layers (use SchNet, etc. when accounting for edge features)
        self.conv1 = GCNConv(hidden_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels)

        # Fully connected layers
        self.lin1 = torch.nn.Linear(hidden_channels, hidden_channels // 2)
        self.lin2 = torch.nn.Linear(hidden_channels // 2, 1)

    def forward(self, x, edge_index, edge_attr, batch):
        # Node embedding
        x = self.embedding(x)
        x = F.relu(x)

        # GCN layers
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.2, training=self.training)

        x = self.conv2(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.2, training=self.training)

        x = self.conv3(x, edge_index)
        x = F.relu(x)

        # Global pooling (aggregate to the crystal level)
        x = global_add_pool(x, batch)

        # Fully connected layers
        x = self.lin1(x)
        x = F.relu(x)
        x = self.lin2(x)

        return x

# Instantiate the model
model_crystal = CGCN(num_node_features=1, hidden_channels=128)

print("===== Crystal Graph Convolutional Network =====")
print(model_crystal)
print(f"\nNumber of parameters: {sum(p.numel() for p in model_crystal.parameters()):,}")

3.4.3 Training Demo with Simulated Data

# Create a simulated dataset (in practice, use Materials Project data)
crystal_dataset = []

for i in range(200):
    num_atoms = torch.randint(4, 12, (1,)).item()
    x = torch.randint(1, 20, (num_atoms, 1)).float()  # atomic numbers

    # Random edges (assumed to be filtered by distance)
    edge_index = torch.randint(0, num_atoms, (2, num_atoms * 4))
    edge_attr = torch.rand(num_atoms * 4, 1) * 5.0  # distance (0-5 Å)

    # Band gap (simulated as a function of atomic number)
    y = (x.mean() / 10.0 + torch.randn(1) * 0.5).clamp(0, 10)

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
    crystal_dataset.append(data)

# Data split
train_crystals = crystal_dataset[:160]
test_crystals = crystal_dataset[160:]

train_loader_crystal = DataLoader(train_crystals, batch_size=16, shuffle=True)
test_loader_crystal = DataLoader(test_crystals, batch_size=16, shuffle=False)

# Training (simplified version)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_crystal = model_crystal.to(device)
optimizer = torch.optim.Adam(model_crystal.parameters(), lr=0.001)
criterion = torch.nn.MSELoss()

print("===== Crystal Band Gap Prediction Training =====")
for epoch in range(1, 51):
    model_crystal.train()
    total_loss = 0

    for data in train_loader_crystal:
        data = data.to(device)
        optimizer.zero_grad()
        out = model_crystal(data.x, data.edge_index, data.edge_attr, data.batch)
        loss = criterion(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs

    if epoch % 10 == 0:
        train_loss = total_loss / len(train_crystals)
        print(f"Epoch {epoch:03d}, Train Loss: {train_loss:.4f}")

# Test evaluation
model_crystal.eval()
test_preds = []
test_targets = []

with torch.no_grad():
    for data in test_loader_crystal:
        data = data.to(device)
        out = model_crystal(data.x, data.edge_index, data.edge_attr, data.batch)
        test_preds.append(out.cpu().numpy())
        test_targets.append(data.y.cpu().numpy())

test_preds = np.concatenate(test_preds)
test_targets = np.concatenate(test_targets)

test_mae = np.mean(np.abs(test_targets - test_preds))
test_r2 = r2_score(test_targets, test_preds)

print(f"\n===== Test Performance =====")
print(f"MAE: {test_mae:.4f} eV")
print(f"R²: {test_r2:.4f}")

3.5 Training Best Practices

3.5.1 Learning Rate Scheduling

from torch.optim.lr_scheduler import ReduceLROnPlateau

# Adjust the learning rate dynamically
scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.5,     # halve the learning rate
    patience=10,    # adjust after 10 epochs without improvement
    verbose=True
)

# Use it inside the training loop
for epoch in range(epochs):
    train_loss = train(model, train_loader, optimizer, criterion, device)
    val_loss = evaluate(model, val_loader, criterion, device)

    # Adjust the learning rate based on the validation loss
    scheduler.step(val_loss)

3.5.2 Early Stopping

class EarlyStopping:
    """
    Early Stopping class
    Stops training when the validation loss stops improving
    """
    def __init__(self, patience=20, min_delta=0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

# Usage example
early_stopping = EarlyStopping(patience=20)

for epoch in range(epochs):
    train_loss = train(model, train_loader, optimizer, criterion, device)
    val_loss = evaluate(model, val_loader, criterion, device)

    early_stopping(val_loss)
    if early_stopping.early_stop:
        print(f"Early stopping at epoch {epoch}")
        break

3.5.3 Data Augmentation (Graph Perturbation)

import torch
from torch_geometric.utils import dropout_edge

def augment_graph(data, drop_edge_prob=0.1, noise_scale=0.01):
    """
    Graph data augmentation

    Parameters:
    -----------
    data : Data
        Original graph
    drop_edge_prob : float
        Probability of dropping an edge
    noise_scale : float
        Scale of the noise added to node features

    Returns:
    --------
    augmented_data : Data
        Augmented graph
    """
    # Edge dropout
    edge_index, edge_mask = dropout_edge(data.edge_index, p=drop_edge_prob)

    # Add noise to node features
    noise = torch.randn_like(data.x) * noise_scale
    x = data.x + noise

    augmented_data = Data(x=x, edge_index=edge_index, y=data.y)
    return augmented_data

# Usage example
original = dataset[0]
augmented = augment_graph(original, drop_edge_prob=0.15)

print(f"Original number of edges: {original.num_edges}")
print(f"Number of edges after augmentation: {augmented.num_edges}")

3.6 Evaluating and Visualizing Model Performance

3.6.1 Computing Evaluation Metrics

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import numpy as np

def evaluate_regression(y_true, y_pred):
    """
    Compute evaluation metrics for a regression model
    """
    mae = mean_absolute_error(y_true, y_pred)
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_true, y_pred)

    # Mean Absolute Percentage Error
    mape = np.mean(np.abs((y_true - y_pred) / y_true)) * 100

    return {
        'MAE': mae,
        'MSE': mse,
        'RMSE': rmse,
        'R²': r2,
        'MAPE': mape
    }

# Usage example
metrics = evaluate_regression(test_targets, test_preds)

print("===== Evaluation Metrics =====")
for name, value in metrics.items():
    print(f"{name}: {value:.4f}")

3.6.2 Residual Plot

import matplotlib.pyplot as plt

def plot_residuals(y_true, y_pred):
    """
    Residual plot
    """
    residuals = y_true - y_pred

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Residuals vs. predicted values
    axes[0].scatter(y_pred, residuals, alpha=0.6, s=20)
    axes[0].axhline(y=0, color='r', linestyle='--', lw=2)
    axes[0].set_xlabel('Predicted value', fontsize=12)
    axes[0].set_ylabel('Residual (actual - predicted)', fontsize=12)
    axes[0].set_title('Residual Plot', fontsize=14)
    axes[0].grid(True, alpha=0.3)

    # Histogram of residuals
    axes[1].hist(residuals, bins=30, alpha=0.7, edgecolor='black')
    axes[1].set_xlabel('Residual', fontsize=12)
    axes[1].set_ylabel('Frequency', fontsize=12)
    axes[1].set_title('Residual Distribution', fontsize=14)
    axes[1].axvline(x=0, color='r', linestyle='--', lw=2)
    axes[1].grid(True, alpha=0.3, axis='y')

    plt.tight_layout()
    plt.show()

# Usage example
plot_residuals(test_targets, test_preds)

3.6.3 Model Comparison

import pandas as pd
import matplotlib.pyplot as plt

# Compare the performance of multiple models
models_performance = {
    'GCN (3 layers)': {'MAE': 0.32, 'R²': 0.88, 'Time': 45.2},
    'GAT (2 layers)': {'MAE': 0.28, 'R²': 0.91, 'Time': 62.8},
    'SchNet': {'MAE': 0.25, 'R²': 0.93, 'Time': 89.5},
    'MPNN': {'MAE': 0.30, 'R²': 0.90, 'Time': 55.1},
}

df = pd.DataFrame(models_performance).T

# Plot
fig, axes = plt.subplots(1, 3, figsize=(16, 4))

# MAE comparison
df['MAE'].plot(kind='bar', ax=axes[0], color='steelblue')
axes[0].set_ylabel('MAE (eV)', fontsize=12)
axes[0].set_title('Mean Absolute Error (lower is better)', fontsize=13)
axes[0].tick_params(axis='x', rotation=45)
axes[0].grid(True, alpha=0.3, axis='y')

# R² comparison
df['R²'].plot(kind='bar', ax=axes[1], color='forestgreen')
axes[1].set_ylabel('R² Score', fontsize=12)
axes[1].set_title('Coefficient of Determination (higher is better)', fontsize=13)
axes[1].tick_params(axis='x', rotation=45)
axes[1].grid(True, alpha=0.3, axis='y')
axes[1].set_ylim(0.8, 1.0)

# Training time comparison
df['Time'].plot(kind='bar', ax=axes[2], color='coral')
axes[2].set_ylabel('Training time (seconds)', fontsize=12)
axes[2].set_title('Computational Cost', fontsize=13)
axes[2].tick_params(axis='x', rotation=45)
axes[2].grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.show()

3.7 Troubleshooting

3.7.1 Common Errors and Solutions

3.7.2 Debugging Checklist

# Check the data
print(f"Number of nodes: {data.num_nodes}")
print(f"Number of edges: {data.num_edges}")
print(f"Isolated nodes: {data.contains_isolated_nodes()}")
print(f"Self-loops: {data.contains_self_loops()}")

# Check tensor shapes
print(f"x.shape: {data.x.shape}")
print(f"edge_index.shape: {data.edge_index.shape}")
print(f"y.shape: {data.y.shape}")

# Check devices
print(f"x device: {data.x.device}")
print(f"edge_index device: {data.edge_index.device}")

# Check the range of the edge index
print(f"max edge index: {data.edge_index.max().item()}")
print(f"num_nodes: {data.num_nodes}")
assert data.edge_index.max().item() < data.num_nodes, "Edge index exceeds the number of nodes"

3.8 Chapter Summary

What We Learned

1. PyTorch Geometric environment setup - Three methods: Conda, pip, and Google Colab - Checking version compatibility and troubleshooting

2. Understanding data structures - Structure of the Data object (x, edge_index, batch) - Graph conversion from RDKit - Batch processing with the DataLoader

3. Hands-on practice with the QM9 dataset - A quantum-chemistry dataset of 134,000 molecules - Implementing and training a GCN model - HOMO-LUMO gap prediction (target of MAE < 0.5 eV)

4. Crystal property prediction - Graph representation of Materials Project crystal data - Crystal Graph Convolutional Network - Band gap prediction

5. Training best practices - Learning rate scheduling - Early Stopping - Data augmentation (graph perturbation)

6. Evaluation and visualization - Metrics such as MAE, MSE, and R² - Residual plots - Performance comparison between models

Key Points

• ✅ PyTorch Geometric is ideal for GNN implementation on materials and molecules
• ✅ QM9 is an excellent molecular property prediction benchmark for beginners
• ✅ Graph preprocessing (checking for isolated nodes and self-loops) is important
• ✅ In batch processing, the batch tensor indicates which graph each node belongs to
• ✅ Learning rate scheduling and Early Stopping prevent overfitting

On to the Next Chapter

In Chapter 4, we will learn advanced GNN techniques: - Graph pooling (hierarchical representations) - Leveraging edge features - Incorporating 3D geometric information (SchNet, DimeNet) - Equivariant GNNs (E(3)-equivariant) - Interpretability with GNNExplainer

Chapter 4: Advanced GNN Techniques →

Exercises

Problem 1 (Difficulty: easy)

List three main attributes contained in a PyTorch Geometric Data object and explain the role of each.

Problem 2 (Difficulty: medium)

A GCN model trained on the QM9 dataset achieved an MAE of 0.8 eV. Propose three concrete approaches to improve its performance.

# Use GAT layers (learn important bonds with the attention mechanism)
from torch_geometric.nn import GATConv

class ImprovedGNN(torch.nn.Module):
    def __init__(self, num_node_features, hidden_channels=128):
        super().__init__()
        # GAT layers (number of heads = 8)
        self.conv1 = GATConv(num_node_features, hidden_channels, heads=8)
        self.conv2 = GATConv(hidden_channels * 8, hidden_channels, heads=8)
        self.conv3 = GATConv(hidden_channels * 8, hidden_channels, heads=1)
        # Increase the number of layers (3 layers → 4 layers)
        self.conv4 = GCNConv(hidden_channels, hidden_channels)

# Incorporate edge features (bond types)
from torch_geometric.nn import NNConv

class EdgeFeaturesGNN(torch.nn.Module):
    def __init__(self, num_node_features, num_edge_features, hidden_channels=64):
        super().__init__()
        # NNConv: accounts for edge features
        nn = torch.nn.Sequential(
            torch.nn.Linear(num_edge_features, hidden_channels * hidden_channels),
            torch.nn.ReLU()
        )
        self.conv1 = NNConv(num_node_features, hidden_channels, nn, aggr='mean')

# Standardize the target variable
y_mean = train_dataset.data.y.mean(dim=0)
y_std = train_dataset.data.y.std(dim=0)

for data in train_dataset:
    data.y = (data.y - y_mean) / y_std

# Data augmentation (graph perturbation)
def augment_graph(data):
    # Edge dropout
    edge_index, _ = dropout_edge(data.edge_index, p=0.1)
    # Add noise
    x = data.x + torch.randn_like(data.x) * 0.01
    return Data(x=x, edge_index=edge_index, y=data.y)

# Double the training data
augmented_train = [augment_graph(data) for data in train_dataset]
train_dataset = train_dataset + augmented_train

Problem 3 (Difficulty: hard)

The following code raised an error. Identify the cause and fix it.

# Code that raises an error
model = GCN_QM9(num_node_features=11, num_classes=1)
device = torch.device('cuda')
model = model.to(device)

for data in train_loader:
    optimizer.zero_grad()
    out = model(data.x, data.edge_index, data.batch)
    loss = criterion(out, data.y)
    loss.backward()
    optimizer.step()

Error message:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

model = GCN_QM9(num_node_features=11, num_classes=1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

for data in train_loader:
    # Move the data to the GPU (important!)
    data = data.to(device)

    optimizer.zero_grad()
    out = model(data.x, data.edge_index, data.batch)
    loss = criterion(out, data.y)
    loss.backward()
    optimizer.step()

# Check the devices
print(f"Model device: {next(model.parameters()).device}")
print(f"Data x device: {data.x.device}")
print(f"Data edge_index device: {data.edge_index.device}")

References

1. Fey, M., & Lenssen, J. E. (2019). "Fast Graph Representation Learning with PyTorch Geometric." ICLR Workshop on Representation Learning on Graphs and Manifolds. GitHub: https://github.com/pyg-team/pytorch_geometric The official PyTorch Geometric paper. Design philosophy and implementation details of the library.

2. Ramakrishnan, R., et al. (2014). "Quantum chemistry structures and properties of 134 kilo molecules." Scientific Data, 1, 140022. DOI: 10.1038/sdata.2014.22 The official QM9 dataset paper. Quantum-chemistry calculation data for 134,000 molecules.

3. Xie, T., & Grossman, J. C. (2018). "Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties." Physical Review Letters, 120(14), 145301. DOI: 10.1103/PhysRevLett.120.145301 The original paper on Crystal Graph Convolutional Networks (CGCN). Application to crystal property prediction.

4. Gilmer, J., et al. (2017). "Neural Message Passing for Quantum Chemistry." ICML 2017. URL: https://arxiv.org/abs/1704.01212 The theory of Message Passing Neural Networks (MPNN). Achieved high-accuracy prediction on QM9.

5. PyTorch Geometric Documentation. (2024). "Introduction by Example." URL: https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html The official PyTorch Geometric tutorial. Illustrates basic usage.

6. RDKit Documentation. (2024). "Getting Started with the RDKit in Python." URL: https://www.rdkit.org/docs/GettingStartedInPython.html The official RDKit documentation. How to create molecule objects from SMILES.

Created: 2025-10-17 Version: 1.0 Template: chapter-template-v2.0 Author: GNN Introduction Series Project