Dive into Graph Neural Networks with PyTorch: A Simple Guide

Abin Varghese

Coffee AI

abin[at]coffeeai[dot]co

June 19, 2024

Abstract

I've been working on Graph Neural Networks lately, and I want to show you how they work. Think of GNNs as neural networks for relationships. Instead of processing images like CNNs, they process networks of interconnected points. In this guide, I'll walk you through building a practical GNN using PyTorch Geometric for classifying research papers in the Cora citation network. You'll see how message passing works and why GNNs are fundamentally different from traditional neural networks.

Introduction

Imagine you have a network of friends. Each friend (node) has certain traits, and the friendships (edges) connect them. A Graph Neural Network helps you understand this network by looking at each friend's traits and how they're connected. Think of it as figuring out who's popular based on who knows whom and their personalities.

GNNs are like Convolutional Neural Networks (CNNs), but instead of working with images, they work with graphs. Networks of interconnected points. They learn by passing messages between nodes, updating their understanding of each node based on its neighbours.

Today, I want to walk you through a project I've been working on. A Graph Neural Network (GNN) implemented in PyTorch for node classification tasks. I'll show you how to build a practical GNN system using the Cora citation network dataset.

The implementation is organized into modular components:

simple_gnn/
├── gnn/
│   ├── __init__.py
│   ├── model.py
│   ├── data.py
│   ├── train.py
│   └── evaluate.py
├── main.py
├── LICENSE
├── README.md
└── requirements.txt

Background: Understanding Graph Neural Networks

Core Mechanism

GNNs operate through three fundamental steps:

Message Passing: Each node gathers information from its neighbours.
Updating Representations: A neural network updates each node's understanding based on the gathered information.
Iteration: This process continues for several rounds to propagate information across the network.

The Cora Citation Network Dataset

I'm using the Cora dataset for this project. It's a collection of scientific publications:

Nodes: Represent research papers (2,708 papers)
Edges: Represent citations between papers (5,429 citations)
Node Features: Bag-of-words representation of paper content (1,433 features)
Classes: Seven categories (e.g., neural networks, rule learning, reinforcement learning)

The goal is to classify papers into their respective categories based on both their content features and citation relationships.

Model Architecture

Components

My GNN model consists of three primary components:

Graph Convolutional Layers (GCNConv): Handle message passing between nodes
ReLU Activation: Introduces non-linearity
Linear Layer: Maps updated node representations to output classes

Implementation Details

The model architecture is defined in gnn/model.py:

gnn/model.py
1import torch
2import torch.nn.functional as F
3from torch_geometric.nn import GCNConv as gcn
4 
5class GNN(torch.nn.Module):
6  def __init__(self, in_channels, hidden_channels, out_channels):
7      super(GNN, self).__init__()
8      self.conv1 = gcn(in_channels, hidden_channels)
9      self.conv2 = gcn(hidden_channels, out_channels)
10      self.linear = torch.nn.Linear(out_channels, out_channels)
11 
12  def forward(self, x, edge_index):
13      x = self.conv1(x, edge_index)
14      x = F.relu(x)
15      x = self.conv2(x, edge_index)
16      x = F.relu(x)
17      x = self.linear(x)
18      return F.log_softmax(x, dim=1)

The GNN class inherits from torch.nn.Module. The __init__ method initializes two graph convolutional layers and a linear layer. The forward method specifies the data flow: input passes through conv1 with ReLU activation, then conv2 with another ReLU, and finally through the linear layer with log-softmax to produce output probabilities. This architecture enables the model to learn by aggregating information from node neighbours.

Training and Evaluation

Data Loading

The gnn/data.py module handles dataset loading:

gnn/data.py
1def load_data(dataset_name='Cora', data_dir='/tmp/Planetoid'):
dataset = Planetoid(root=data_dir, name=dataset_name)
data = dataset[0]
train_mask = data.train_mask
val_mask = data.val_mask
test_mask = data.test_mask

return data, train_mask, val_mask, test_mask

This function fetches the specified dataset (default is Cora) from the given directory, retrieving the data and masks for training, validation, and testing. This simplifies data preparation for the model.

Training Process

The gnn/train.py module defines the training loop:

gnn/train.py
1def train_model(model, data, train_mask, val_mask, optimizer, criterion, num_epochs=200):
model.train()
3 
for epoch in range(num_epochs):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[train_mask], data.y[train_mask])
    loss.backward()
    optimizer.step()
10 
    if epoch % 10 == 0:
        model.eval()
        with torch.no_grad():
            out = model(data.x, data.edge_index)
            val_acc = torch.sum(out[val_mask].argmax(dim=1) == data.y[val_mask]) / val_mask.sum()
            print(f'Epoch: {epoch}, Validation Accuracy: {val_acc:.4f}')

return model

Each epoch zeros out gradients, performs a forward pass to compute output and loss, backpropagates to update weights, and periodically evaluates validation accuracy.

Model Evaluation

The gnn/evaluate.py module assesses model performance:

gnn/evaluate.py
1def evaluate_model(model, data, test_mask):
model.eval()
with torch.no_grad():
    out = model(data.x, data.edge_index)
    test_acc = torch.sum(out[test_mask].argmax(dim=1) == data.y[test_mask]) / test_mask.sum()
    print(f'Test Accuracy: {test_acc:.4f}')

This sets the model to evaluation mode, computes predictions without updating weights, and calculates test accuracy by comparing predictions to actual labels.

Complete Pipeline

The main.py script orchestrates the entire workflow:

main.py
1def main():
2  # Load data
3  data, train_mask, val_mask, test_mask = load_data()
4 
5  # Get the number of classes
6  num_classes = data.y.max().item() + 1
7 
8  # Define model, optimizer, and loss function
9  model = GNN(data.num_node_features, 16, num_classes)
10  optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
11  criterion = torch.nn.CrossEntropyLoss()
12 
13  # Train the model
14  trained_model = train_model(model, data, train_mask, val_mask, optimizer, criterion)
15 
16  # Evaluate the model
17  evaluate_model(trained_model, data, test_mask)
18 
19if __name__ == "__main__":
20  main()

This file ties together data loading, model setup, training, and evaluation in a cohesive pipeline.

Running the Implementation

Setup on Windows

Windows Setup
1git clone https://github.com/Spartan-119/simple_gnn.git
2cd simple_gnn
3python -m venv gnn_venv
4gnn_venv\Scripts\activate
5pip install -r requirements.txt
6python main.py

Setup on MacOS/Linux

MacOS/Linux Setup
1git clone https://github.com/Spartan-119/simple_gnn.git
2cd simple_gnn
3python3 -m venv gnn_venv
4source gnn_venv/bin/activate
5pip install -r requirements.txt
6python3 main.py

This downloads the Cora dataset, trains the GNN, and evaluates its performance.

Training Output — **Figure:** Training progress showing validation accuracy over epochs

Discussion: GNNs vs CNNs

Fundamental Differences

CNNs excel at grid-like data such as images, but they fall short with graphs where connections are irregular and varied. GNNs address this limitation:

Flexibility: Can work with any graph structure, not just grids
Message Passing: Nodes update their state by aggregating information from neighbours, unlike CNNs which use fixed filters over fixed positions
Dynamic: Can adapt to different graph sizes and shapes, whereas CNNs have fixed input sizes

Key Advantages

I've found that GNNs provide several advantages over traditional neural networks:

Capture Relationships: Understand not just data points but how they're connected
Generalisation: Apply to various types of graphs, whether social networks, molecules, or transport systems
Permutation Invariance: Output is independent of node ordering

Applications

GNNs have broad applications across domains:

Social Networks: Analyzing user connections to recommend friends or detect communities
Biological Networks: Understanding protein interactions or gene regulatory networks
Recommendation Systems: Suggesting products based on user-item interaction graphs
Traffic Networks: Predicting traffic flow by analyzing road networks
Knowledge Graphs: Enhancing search engines by connecting related concepts

Customization and Extension

Model Customization

You can modify the GNN by editing:

model.py: Change the architecture (add layers, adjust dimensions)
data.py: Load different datasets or preprocess data differently
main.py: Modify hyperparameters like learning rate or hidden dimensions

Future Directions

The field of GNNs is rapidly evolving. Exciting future directions include:

Scalability: Making GNNs handle larger graphs efficiently (millions of nodes)
Dynamic Graphs: Dealing with graphs that change over time
Interpretability: Understanding how GNNs make decisions, crucial for fields like medicine
Heterogeneous Graphs: Handling graphs with different types of nodes and edges
Graph Generation: Creating new graph structures for drug discovery or molecular design

Conclusion

I've shown you a practical implementation of Graph Neural Networks using PyTorch Geometric. By structuring the code into modular components, I've created a flexible framework for node classification tasks. The Cora dataset serves as an excellent starting point for understanding how GNNs leverage both node features and graph structure to make predictions.

GNNs represent a powerful paradigm shift in deep learning. They extend neural networks beyond Euclidean data to the rich world of relational structures. With this foundation, you can start experimenting with GNNs and exploring the fascinating world of graph-structured data.

The complete source code is available at github.com/Spartan-119/simple_gnn.