Getting Started with Graph Neural Networks in Python

Graph neural networks (GNNs) have emerged as a powerful approach for machine learning on graph-structured data. Graphs are a natural way to represent many real-world datasets where entities are related to each other in complex ways, such as social networks, molecular structures, power grids, web networks, and knowledge graphs. The key idea behind GNNs is to learn embeddings for nodes in the graph that capture both the node features and graph structure, and use these embeddings for downstream machine learning tasks like node classification, link prediction, and graph classification.

In recent years, GNNs have achieved state-of-the-art performance on many benchmark graph datasets and are being increasingly used for applications like drug discovery, traffic forecasting, recommender systems, and fake news detection. The field is rapidly growing with new GNN architectures and techniques being proposed at a fast pace.

The goal of this blog post is to provide a comprehensive beginner‘s guide to getting started with GNNs in Python. We‘ll cover the key concepts, walk through hands-on code examples, and discuss practical considerations and challenges. By the end of this post, you‘ll be equipped with the knowledge and tools to start applying GNNs to your own graph-based machine learning problems.

Background on Graphs and Graph Neural Networks

Before we dive into graph neural networks, let‘s start with a quick recap of graphs. A graph $G=(V,E)$ is a data structure consisting of a set of nodes (or vertices) $V$, and a set of edges $E$ that connect pairs of nodes. Edges can be directed or undirected, and can optionally have associated weights. The adjacency matrix $A$ is a square matrix where $A_{ij} = 1$ if there is an edge from node $i$ to node $j$, and 0 otherwise.

Traditionally, machine learning on graphs has relied on hand-engineered features like graph statistics, kernel functions, or unsupervised node embeddings. However, these approaches often fail to capture the complex dependencies and long-range interactions between nodes.

Graph neural networks aim to address this limitation by learning node representations that directly incorporate both node features and the graph structure. The key idea is to iteratively update the embedding of each node by aggregating information from its neighbors. Stacking multiple GNN layers allows information to propagate through the graph, so that the final node embeddings capture long-range dependencies. These node embeddings can then be used for downstream machine learning tasks.

Some of the pioneering works that established GNNs include:

• Gori et al. (2005) proposed recursive neural networks for learning on directed acyclic graphs
• Scarselli et al. (2009) introduced graph neural networks as a general framework for neural networks on graphs
• Bruna et al. (2013) developed spectral convolutional neural networks based on the graph Laplacian
• Duvenaud et al. (2015) proposed convolutional networks for learning molecular fingerprints
• Kipf & Welling (2017) popularized graph convolutional networks (GCNs) for semi-supervised node classification

Since then, there has been an explosion of research on new GNN architectures, extensions to more complex graph types, scalable training techniques, and applications to a wide variety of domains. For an extensive survey of GNNs, see Zhou et al. (2018) or Wu et al. (2021).

Components of a Graph Neural Network

A graph neural network typically consists of the following components:

1. Input graph: The input to a GNN is a graph $G=(V,E)$ with node features $X \in \mathbb{R}^{N \times F}$ where $N$ is the number of nodes and $F$ is the number of input features per node.

2. GNN layers: The core of a GNN is a stack of neural network layers that operate on the graph structure to update the node embeddings. Each layer performs some form of message passing where nodes aggregate information from their neighbors to compute new embeddings. Different GNN architectures vary in how the aggregation is performed, e.g., using sum, mean, max-pooling, or attention mechanisms.

3. Readout: To generate graph-level predictions, a readout function is used to aggregate the final node embeddings into a graph embedding. This can be a simple sum or average of node embeddings, or more sophisticated pooling methods.

4. Loss function: GNNs are typically trained end-to-end using standard loss functions depending on the downstream task, such as cross-entropy for node classification or binary cross-entropy for link prediction.

During message passing, a GNN layer updates the embedding $h_i$ of each node $i$ by aggregating the embeddings of its neighbors ${h_j : j \in N(i)}$ where $N(i)$ denotes the set of neighbors of node $i$. A common form of aggregation is:

$h_i^{(l+1)} = \sigma(W^{(l)} \cdot \text{AGG}({h_j^{(l)} : j \in N(i)}) + b^{(l)})$

where $h_i^{(l)}$ is the embedding of node $i$ at layer $l$, $W^{(l)}$ and $b^{(l)}$ are learnable weights and biases, $\text{AGG}$ is an aggregation function, and $\sigma$ is a non-linear activation like ReLU.

By stacking multiple GNN layers, information can propagate through the graph, allowing the model to capture long-range interactions between nodes. The final node embeddings can be used for node-level tasks like node classification, while the graph embedding from the readout can be used for graph-level tasks.

Implementing Graph Neural Networks in Python

There are several deep learning libraries that provide built-in support for graph neural networks in Python, including:

• PyTorch Geometric: A library built on top of PyTorch for deep learning on graphs
• Deep Graph Library (DGL): A Python package for deep learning on graphs
• StellarGraph: A Python library for machine learning on graphs, focused on tasks like node classification, link prediction, graph classification, etc.
• TensorFlow GNN: TensorFlow‘s library for building graph neural networks

For this post, we‘ll focus on using PyTorch Geometric (PyG), as it has a simple and intuitive API, and supports many state-of-the-art GNN architectures. To install PyG, you can use pip:

pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric

PyG uses a custom Data object to represent a graph, which stores the node features, edge indices, edge features, and optional node or graph labels. Here‘s an example of loading the Cora citation network dataset:

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root=‘data/Planetoid‘, name=‘Cora‘)

This loads the Cora dataset, which consists of 2708 scientific publications classified into one of seven classes. Each publication is represented by a 1433-dimensional binary bag-of-words feature vector, and citations between papers are represented as undirected edges.

To define a GNN model in PyG, we create a new class that inherits from torch.nn.Module and build the model using the provided GNN layers. Here‘s an example of a simple 2-layer Graph Convolutional Network (GCN) for node classification:

import torch
from torch.nn import Linear
from torch_geometric.nn import GCNConv

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

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

The GCNConv layer performs message passing by aggregating neighboring node embeddings weighted by normalized edge weights. We can instantiate the model and train it for node classification as follows:

model = GCN(num_features=dataset.num_features, hidden_channels=16, num_classes=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train(data):
    model.train()
    optimizer.zero_grad()  
    out = model(data.x, data.edge_index)  
    loss = torch.nn.CrossEntropyLoss()(out[data.train_mask], data.y[data.train_mask])  
    loss.backward()  
    optimizer.step()  
    return loss

for epoch in range(200):
    loss = train(dataset[0])
    print(f‘Epoch: {epoch:03d}, Loss: {loss:.4f}‘)

After training, we can evaluate the model‘s performance on the test set:

model.eval()
_, pred = model(dataset[0].x, dataset[0].edge_index).max(dim=1)
correct = float(pred[dataset[0].test_mask].eq(dataset[0].y[dataset[0].test_mask]).sum().item())
acc = correct / dataset[0].test_mask.sum().item()
print(f‘Test Accuracy: {acc:.4f}‘)

This simple GCN model achieves around 81% test accuracy on the Cora dataset, comparable to the results reported in the original GCN paper.

Advanced Topics and Extensions

There are many ways to extend and improve upon the basic GCN model we covered:

• More sophisticated GNN layers: There are many other types of GNN layers that can capture different aspects of the graph structure, such as Graph Attention Networks (GATs) which use self-attention to weight the importance of neighboring nodes, Graph Isomorphism Networks (GINs) which are provably as powerful as the Weisfeiler-Lehman graph isomorphism test, and Graph Sample and Aggregate (GraphSAGE) which is designed for inductive learning on large graphs.

• Heterogeneous and dynamic graphs: Many real-world graphs are heterogeneous, meaning they have multiple types of nodes and edges, or dynamic, meaning the nodes and edges can change over time. Extensions of GNNs like Relational GCNs and Temporal GNNs can handle these more complex graph types.

• Pre-training graph neural networks: Similar to pre-training in natural language processing, pre-training GNNs on large unlabeled graphs using self-supervised learning has been shown to improve performance on downstream tasks. Some popular pre-training methods include node attribute masking, edge prediction, and graph-level contrastive learning.

• Scalability: Training GNNs on large graphs can be challenging due to memory constraints and computational complexity. Techniques like graph sampling, subgraph batching, and layer-wise sampling can help scale GNNs to graphs with millions of nodes and edges.

Challenges and Open Problems

Despite the impressive progress made in graph neural networks, there are still many challenges and open problems in the field:

• Over-smoothing and vanishing gradients: As the number of GNN layers increases, the node embeddings can become oversmoothed and lose discriminative power. This can also lead to vanishing gradients during training, making it difficult to train very deep GNNs.

• Extending to inductive settings: Most GNNs are transductive, meaning they require all nodes to be present during training and can only make predictions for those nodes. Extending GNNs to inductive settings where the model needs to generalize to entirely unseen graphs is an active area of research.

• Incorporating edge features and attributes: While most GNNs focus on learning node embeddings, many graphs also have rich edge features that could be informative for downstream tasks. Developing GNN architectures that can effectively incorporate edge features is an important direction.

• Interpretability: Like many deep learning models, GNNs can be difficult to interpret and explain. Developing techniques to understand what GNNs are learning and how they arrive at their predictions is crucial for building trust in GNN models and applying them to high-stakes domains.

Conclusion

In this post, we covered the basics of graph neural networks and how to implement them in Python using PyTorch Geometric. We saw how GNNs leverage both node features and graph structure to learn powerful node and graph embeddings, and how they can be applied to tasks like node classification, link prediction, and graph classification.

GNNs are a rapidly growing field with many exciting developments in recent years. From molecular property prediction to fake news detection to traffic forecasting, GNNs are being applied to an increasingly wide range of domains. At the same time, there are still many challenges and open problems to be addressed, from scalability to interpretability to inductive learning.

I encourage you to try implementing GNNs on your own graph datasets and experiment with different architectures and techniques. The PyTorch Geometric documentation and examples are a great resource to learn more. I hope this post has given you a solid foundation to start your journey with graph neural networks. Happy learning!