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.Anatomy of a PyG Training ScriptA 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.digraph G { rankdir=TB; node [shape=box, style="rounded,filled", fillcolor="#e9ecef", fontname="sans-serif"]; edge [fontname="sans-serif"]; subgraph cluster_data { label = "1. Data Preparation"; style=filled; color="#f8f9fa"; LoadData [label="Load PyG Dataset (e.g., Cora)"]; DataObject [label="Access Graph Data\n(data.x, data.edge_index, data.y, masks)"]; LoadData -> DataObject; } subgraph cluster_model { label = "2. Model Definition"; style=filled; color="#f8f9fa"; DefineModel [label="Define GNN Architecture\n(inherits from torch.nn.Module)"]; InstantiateModel [label="Instantiate Model & Optimizer"]; DefineModel -> InstantiateModel; } subgraph cluster_training { label = "3. Training & Evaluation"; style=filled; color="#f8f9fa"; TrainLoop [label="For each Epoch:\n- Train Function\n- Evaluate Function"]; TrainFunc [label="train()\nForward Pass\nCompute Loss (on train nodes)\nBackpropagation"]; EvalFunc [label="evaluate()\nForward Pass\nCompute Accuracy (on val/test nodes)"]; TrainLoop -> TrainFunc [label="call"]; TrainLoop -> EvalFunc [label="call"]; } DataObject -> InstantiateModel [style=dashed, label="num_features, num_classes"]; InstantiateModel -> TrainFunc [style=dashed]; DataObject -> TrainFunc [style=dashed, label="pass data"]; InstantiateModel -> EvalFunc [style=dashed]; DataObject -> EvalFunc [style=dashed, label="pass data"]; }The workflow of a standard PyTorch Geometric training script, from data loading to model evaluation.1. Imports and Data LoadingFirst, 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.2. Defining the GNN ModelNext, 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.3. The Training LoopThe 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.4. The Evaluation FunctionSimilarly, 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 accsIn 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.5. Putting It All TogetherFinally, 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.{"layout":{"xaxis":{"title":"Epoch"},"yaxis":{"title":"Value"},"title":"GNN Training Progress","legend":{"traceorder":"normal"}},"data":[{"type":"scatter","name":"Loss","x":[10,20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200],"y":[1.75,1.43,1.15,0.92,0.76,0.65,0.58,0.53,0.49,0.46,0.44,0.42,0.41,0.40,0.39,0.38,0.38,0.37,0.37,0.36],"line":{"color":"#f03e3e"}},{"type":"scatter","name":"Validation Accuracy","x":[10,20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180,190,200],"y":[0.55,0.68,0.74,0.77,0.78,0.79,0.79,0.80,0.80,0.81,0.81,0.81,0.81,0.81,0.81,0.81,0.82,0.82,0.82,0.82],"yaxis":"y2","line":{"color":"#228be6"}}],"frames":[],"layout_yaxis2":{"title":"Accuracy","overlaying":"y","side":"right"}}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.