A complete training pipeline for Graph Neural Networks (GNNs) can be constructed using PyTorch Geometric components, such as the Data object and modular GNN layers. A single, reusable script is built for training a GNN on a node classification task. The resulting structure provides a standard pattern adaptable for many graph-based projects.
We will walk through each logical block of the script: loading the data, defining the model, creating the training and evaluation loops, and finally, orchestrating the entire process.
A typical training script in PyTorch Geometric follows a well-defined structure that integrates data handling, model definition, and the training process. This separation of concerns makes the code clean, readable, and easy to modify.
The workflow of a standard PyTorch Geometric training script, from data loading to model evaluation.
First, we import the necessary libraries. This includes torch for core tensor operations, torch.nn.functional for activation functions and loss functions, and several components from torch_geometric. We'll use the Planetoid dataset collection, which includes the Cora, CiteSeer, and PubMed benchmark datasets.
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
# Load the Cora dataset
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
print(f'Dataset: {dataset.name}')
print(f'Number of graphs: {len(dataset)}')
print(f'Number of nodes: {data.num_nodes}')
print(f'Number of edges: {data.num_edges}')
print(f'Number of features: {dataset.num_node_features}')
print(f'Number of classes: {dataset.num_classes}')
The Planetoid object downloads the data and processes it into a format PyG can use. Since Cora is a single graph, the dataset contains just one Data object, which we access with dataset[0]. This data object conveniently holds not just the node features (data.x) and graph connectivity (data.edge_index), but also the ground-truth labels (data.y) and the predefined masks (data.train_mask, data.val_mask, data.test_mask) for semi-supervised learning.
Next, we define our network architecture. We create a class that inherits from torch.nn.Module, which is the standard way to build models in PyTorch. This allows us to use PyG layers just like any other PyTorch layer.
Here, we'll define a simple two-layer GCN. The first GCNConv layer maps the input node features to a hidden dimension, and the second layer maps the hidden representations to the number of output classes.
class GCN(torch.nn.Module):
def __init__(self, num_features, num_classes):
super(GCN, self).__init__()
self.conv1 = GCNConv(num_features, 16)
self.conv2 = GCNConv(16, num_classes)
def forward(self, x, edge_index):
# First GCN layer
x = self.conv1(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=0.5, training=self.training)
# Second GCN layer
x = self.conv2(x, edge_index)
# Apply LogSoftmax for the NLLLoss function
return F.log_softmax(x, dim=1)
The forward method defines how data flows through the network. It takes the node feature matrix x and the edge_index as input. Notice the use of a ReLU activation function and dropout for regularization after the first layer. The final output is passed through a log_softmax function, which is often paired with the Negative Log Likelihood loss (F.nll_loss) for classification tasks.
The training loop contains the logic for optimizing the model's parameters. We define a function that performs a single pass of training.
def train(model, optimizer, data):
model.train() # Set the model to training mode
optimizer.zero_grad() # Clear gradients
# Perform a single forward pass
out = model(data.x, data.edge_index)
# Compute the loss using only the training nodes
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
# Derive gradients
loss.backward()
# Update parameters
optimizer.step()
return loss.item()
This function encapsulates the standard PyTorch training steps. The most important part for GNNs is how the loss is calculated. Instead of using the output for all nodes, we apply data.train_mask to select only the logits and labels for the nodes in the training set. This is the essence of semi-supervised node classification.
Similarly, we create a function to evaluate the model's performance on the validation or test set. This function should not update the model's weights, so we operate under torch.no_grad().
@torch.no_grad()
def test(model, data):
model.eval() # Set the model to evaluation mode
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1) # Use the class with highest probability
# Check against the validation and test sets
accs = []
for mask in [data.train_mask, data.val_mask, data.test_mask]:
correct = pred[mask] == data.y[mask] # Check against ground-truth labels
accs.append(int(correct.sum()) / int(mask.sum())) # Calculate accuracy
return accs
In test(), we set the model to eval() mode, which disables layers like dropout. We get predictions for all nodes by taking the argmax of the output logits. Then, we compute the accuracy for each data split (train, validation, and test) by applying the corresponding masks.
Finally, we write the main script logic that initializes everything and runs the training process. We instantiate the model, define an optimizer like Adam, and then loop for a set number of epochs. In each epoch, we call our train and test functions and log the results.
# Instantiate the model
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)
# Define the optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
# Training execution
best_val_acc = 0
for epoch in range(1, 201):
loss = train(model, optimizer, data)
train_acc, val_acc, test_acc = test(model, data)
# Simple early stopping
if val_acc > best_val_acc:
best_val_acc = val_acc
if epoch % 10 == 0:
print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, '
f'Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}, Test Acc: {test_acc:.4f}')
This main block drives the learning process. By tracking the validation accuracy, you can implement simple forms of early stopping or save the model that performs best on the validation set. After 200 epochs, the model trained on the small Cora dataset typically achieves high accuracy on the test nodes.
Typical training curves for a GCN on the Cora dataset. The training loss decreases while the validation accuracy increases and stabilizes.
With this complete script, you have a powerful and flexible template for training GNNs. The final hands-on section will challenge you to apply this structure to train and evaluate a model on the Cora dataset from start to finish.
Was this section helpful?
© 2026 ApX Machine LearningEngineered with