Sampling methods for Graph Neural Networks

The versatility of graph representations makes them highly valuable for solving a wide range of problems, each with its own unique data structure. However, generating universal embeddings that apply across different applications remains a significant challenge. 

This article explores PyTorch Geometric (PyG) data loaders and presents a streamlined interface to abstract away their underlying complexity

Overview
Graph Neural Networks
PyTorch Geometric
Graph Samplers
       Overview
       Data Splits
Common Sampling Architectures
       Random Node Sampling
       Neighbor Node Sampling
       Neighbor Link Sampling
       Subgraphs Cluster Sampling
       Graph Sampling Based Inductive Learning Method
Implementation
       Environment
       Datasets
       Graph Data Loader
       Random Node Loader
       Neighbor Node Loader
References        

What you will learn: how graph data loaders influence node classification in a Graph Neural Network implemented with PyTorch Geometric and leverage easy-to-use interfaces to kickstart your exploration of the various architectures.

✅ The complete article, featuring design principles, detailed implementation, in-depth analysis, and exercises, is available at Demystifying Graph Sampling & Walk Methods

Overview

The wide range of datasets and samplers available in PyTorch Geometric can be overwhelming at first. To simplify exploration, we’ve developed interfaces that make it easier to navigate and experiment with different datasets and sampling strategies.

Data on manifolds can often be represented as a graph, where the manifold's local structure is approximated by connections between nearby points. GNNs and their variants (like Graph Convolutional Networks (GCNs)) extend neural networks to process data on non-Euclidean domains by leveraging the graph structure, which may approximate the underlying manifold. Application Social network analysis, molecular structure prediction, and 3D point cloud data can all be modeled using GNNs.

I strongly recommend to review my article introducing PyTorch Geometric Taming PyTorch Geometric for Graph Neural Networks

📌 The description of the inner-workings and mathematical foundation of graph neural networks, message passing architecture and aggregation policies are beyond the scope of this article. Some of the most relevant tutorials and presentations on GNNs are listed in references [ref 2, 3, 4 & 5].

Graph Neural Networks

⚠️This article explores the PyG (PyTorch Geometric) Python library to evaluate various graph neural network (GNN) architectures. It is not intended as an introduction or overview of GNNs and assumes the reader has some prior knowledge of the subject.

The variety of Graph Neural Network (GNN) architectures has grown rapidly over the past three years, driven by a surge in technical publications on the subject.

Here’s a shortlist of some of the most widely used GNN models.

• Graph Convolutional Networks (GCNs): GCNs generalize the concept of convolution from grids (e.g., images) to graphs. They aggregate information from a node's neighbors using normalized adjacency matrices and apply transformations to learn node embeddings.
• Graph Attention Networks (GATs): GATs use attention mechanisms to learn the importance of neighboring nodes dynamically. Each edge is assigned a learned weight during aggregation.
• GraphSAGE (Graph Sample and Aggregate): It learns node embeddings by sampling and aggregating features from a fixed-size neighborhood of each node, enabling scalable learning on large graphs.
• Graph Isomorphism Networks (GINs): GINs are designed to be as powerful as the Weisfeiler-Lehman (WL) graph isomorphism test, distinguishing graph structures more effectively.
• Spectral Graph Neural Networks (SGNN): These networks operate in the spectral domain using the graph Laplacian. They use eigenvectors of the Laplacian for convolution-like operations.
• Graph Pooling Networks: They summarize graph information into a smaller representation, similar to pooling in CNNs. They can be categorized into Global and hierarchical pooling
• Hyperbolic Graph Neural Networks: These networks operate in hyperbolic space, which is well-suited for representing hierarchical or tree-like graph structures.
• Dynamic Graph Neural Networks: These networks are designed to handle temporal graphs, where nodes and edges evolve over time.
• Relational Graph Convolutional Networks (R-GCNs): R-GCNs extend GCNs to handle heterogeneous graphs with different types of nodes and edges.
• Graph Transformers: They adapt the Transformer architecture to graph-structured data using attention mechanisms and global context.
• Graph Autoencoders: These are used for unsupervised learning on graphs, aiming to reconstruct graph structure and node features.
• Diffusion-Based GNNs: These networks use graph diffusion processes to propagate information.

PyTorch Geometric

PyTorch Geometric (PyG) is a graph deep learning library built on PyTorch, designed for efficient processing of graph-structured data. It provides essential tools for building, training, and deploying Graph Neural Networks (GNNs) [ref 6].

The key Features of PyG are:

• Efficient Graph Processing to optimize memory and computation using sparse graph representations.
• Flexible GNN Layers that includes GCN, GAT, GraphSAGE, GIN, and other advanced architectures.
• Batching for Large Graphs to support mini-batching for handling graphs with millions of edges.
• Seamless PyTorch Integration with full compatibility with PyTorch tensors, autograd, and neural network modules.
• Diverse Graph Support for directed, undirected, weighted, and heterogeneous graphs.

The most important PyG Modules are:

• torch_geometric.data to manages graph structures, including nodes, edges, and features.
• torch_geometric.nn to provide data scientists prebuilt GNN layers like convolutional and gated layers.
• torch_geometric.transforms to pre-process input data (e.g., feature normalization, graph sampling).
• torch_geometric.loader to handle large-scale graph datasets with specialized loaders.

Graph Samplers

Overview

Some real-world applications involve handling extremely large graphs with thousands of nodes and millions of edges, posing significant challenges for both machine learning algorithms and visualization. Fortunately, PyG (PyTorch Geometric) enables data scientists to batch nodes or edges, effectively reducing computational overhead for training and inference in graph-based models.

First we need to introduce the attributes of the data of type torch_geometric.data.Data that underline the representation of a graph in PyG.

Data Splits

The graph is divided into training, validation, and test datasets by assigning train_mask, val_mask, and test_mask attributes to the original Data object, as demonstrated in the following code snippet.

# 1. Define the indices for training, validation and test data points
train_idx = torch.tensor([0, 1, 2, 4, 6, 7, 8, 11, 12, 13, 14]) 
val_idx = torch.tensor([3, 9, 14])
test_idx = torch.tensor([5, 10])

#2. Verify all indices are accounted for with no overlap
validate_split(train_idx, val_idx, test_idx)

#3. Get the training, validation and test data set
train_data = data.x[train_idx], data.y[train_idx]
val_data = data.x[val_idx], data.y[val_idx]
test_data = data.x[test_idx], data.y[test_idx]        

Alternatively, we can use the RandomNodeSplit and RandomLinkSplit classes to directly extract the training, validation, and test datasets.

from torch_geometric.transforms import RandomNodeSplit

transform = RandomNodeSplit(is_undirected=True)
train_data, val_data, test_data = transform(data)        

Common Sampling Architectures

The graph nodes and link samplers are an extension of PyTorch ubiquitous data loader. A node loader performs a mini-batch sampling from node information and a link loader performs a similar mini-batch sampling from link information.'

The latest version of PyG supports an extensive range of graph data loaders. Below is an illustration of the most commonly used node and link loaders.

📌 As PyG implements node and edge sampling methods through data set loaders in PyG, we use the terms sampler and loader interchangeably.

Random Node Sampling

A data loader that randomly samples nodes from a graph and returns their induced subgraph. In this case, the two sampled subgraphs are highlighted in blue and red.

Neighbor Node Sampling

This loader partitions nodes into batches and expands the subgraph by including neighboring nodes at each step. Each batch, representing an induced subgraph, starts with a root node and attaches a specified number of its neighbors. This approach is similar to breath-first search in trees.

Neighbor Link Sampling

This loader is similar to the neighborhood node loader. It partitions links and associated nodes into batches and expands the subgraph by including neighboring nodes at each step.

Subgraphs Cluster Sampling

Divides a graph data object into multiple subgraphs or partitions. A batch is then formed by combining a specified number (`batch_size`) of subgraphs. In this example, two subgraphs, each containing five green nodes, are grouped into a single batch.

Graph Sampling Based Inductive Learning Method

This is an inductive learning approach that enhances training efficiency and accuracy by constructing mini-batches through sampling subgraphs from the training graph, rather than selecting individual nodes or edges from the entire graph. This approach is similar to depth-first search in trees.

Implementation

Environment

• Source code is available at Github.com/patnicolas/geometriclearning/dataset/graph
• To enhance the readability of the algorithm implementations, we have omitted non-essential code elements like error checking, comments, exceptions, validation of class and method arguments, scoping qualifiers, and import statements.

⚠️ Warning: Some sampling methods in PyTorch Geometric rely on additional modules: torch-sparse, torch-scatter, torch-spline-conv, and torch-cluster. These are dependencies of torch-geometric but they may not be compatible with the latest versions of PyTorch, particularly across different operating systems.

For macOS, we recommend the following version setup for best compatibility:

• Python: 3.11.8
• PyTorch: 2.1.0
• torch-geometric: 2.6.1
• torch-sparse: 0.6.18
• torch-scatter: 2.1.2
• torch-spline-conv: 1.2.2
• torch-cluster: 1.6.3

These modules currently support only CPU and CUDA execution — MPS (Metal) is not supported.

Datasets

The PyTorch Geometric library offers a wide range of datasets that data scientists can use to evaluate different Graph Neural Network architectures [ref 7].

Since each dataset may follow its own loading and extraction protocol, we'll create a convenient wrapper class called, PyGDatasets to serve as a unified interface—abstracting away the underlying implementation details.

Here is an example of a private loading method:

The data for each source can be simply extracted as:

CiteSeer:
Data(x=[3327, 3703], edge_index=[2, 9104], y=[3327], train_mask=[3327], val_mask=[3327], test_mask=[3327])        

Graph Data Loader

Let's analyze the impact of different graph data loaders on the performance of a Graph Convolutional Neural Network (GCN).

To facilitate this evaluation, we'll create a wrapper class, GraphDataLoader, for managing the access and loading of graph data. 

As shown in code snippet #4, we've organized the different loading and sampling methods into a dictionary called loader_sampler_dict, where each sampler name serves as a key and the corresponding data loading lambda function is the value.

The arguments of the constructor are:

• loader_attributes: Dictionary for the configuration of this specific loader
• dataset_name: Name of the PyG dataset

The dictionary of the attributes (1) is specific to the loader and therefore needs to be validated (2). For instance the attributes for accessing Flickr dataset using neighbor loader is defined as

loader_attributes={
   'id': 'NeighborLoader', 
   'num_neighbors': [3, 2], 
   'batch_size': 4, 
   'replace': True, 
   'num_workers': 1
}        

The data is extracted as the first entry in the dataset (3).

The call method directs requests to the appropriate node or link loader-sampler.

To keep this article concise, our evaluation focuses on the following two graph data loaders:

• Random Nodes
• Neighbors Nodes

Random Node Loader

The only configuration 3 attributes for the random node loader are 

• num_parts: Number of subgraphs generated through partitioning and sampling of the input graph. The data set is split into num_parts subgraphs to improve performance and parallelization for large graphs. 
• batch_size: Number of graph nodes used in a mini-batch
• num_workers: Number of concurrent execution for training and validation.

📌 The meaning of batch_size varies depending on the type of graph task: for graph-level tasks (e.g., graph classification), it refers to the number of graphs per mini-batch; for node-level tasks, it’s the number of nodes per mini-batch; and for edge-level tasks, it represents the number of edges per mini-batch.

The loader for the training set shuffles the data while the order of data points for the validation set is preserved.

We consider the Flickr data set included in PyTorch Geometric. The Flickr dataset is a graph where nodes represent images and edges signify similarities between them [ref 8]. It includes 89,250 images and 899,756 relationships. Node features consist of image descriptions and shared properties.

The purpose is to classify Flickr images (defined as graph nodes) into one of the 108 categories.

0: Data(x=[349, 500], edge_index=[2, 8], y=[349], train_mask=[349], val_mask=[349], test_mask=[349])
1: Data(x=[349, 500], edge_index=[2, 12], y=[349], train_mask=[349], val_mask=[349], test_mask=[349])
2: Data(x=[349, 500], edge_index=[2, 6], y=[349], train_mask=[349], val_mask=[349], test_mask=[349])        

We train a three-layer Graph Convolutional Neural Network (GCN) on the Flickr dataset for classify these images into 108 categories. In this first experiment the model trains on data extracted by the random node loader. For clarity, the code for training the model, computing losses, and evaluating performance metrics has been omitted.

📌 The training and evaluation of Graph Neural Networks is outside the scope of this article and will be covered in future installments.

The following plot tracks the various performance metrics (accuracy, precision and recall) as well as the training and validation loss over 60 iterations.

Neighbor Node Loader

The configuration parameters used for this loader include:

• num_neighbors: Number of neighbors to sample at each layer or hop in a Graph Neural Network. It is defined as an array, e.g., [num_neighbors_first_hop, num_neighbors_second_hop, ...].
• replace: Flag that specifies if a sampling is performed with or without replacement.
• batch_size: Number of graph nodes in a min_batch.
• num_workers: Number of concurrent executions for training and validation

We specify few other parameters which value do not vary during our evaluation>

• drop_last: Flat that specifies whether the last min-batch has to be dropped if its size < batch_size
• input_nodes Target graph nodes (training or validation) using data split and mask.

📌 Increasing the number of hops when sampling neighboring nodes helps capture long-range dependencies, but it comes with increased computational cost. Moreover, for most tasks, performance tends to degrade beyond three hops. Regardless of the depth, the number of sampled nodes in the first hop should be larger than in the second hop, and this pattern should continue for subsequent hops

Running the code described in snippet 7, with the attributes

{
  'id': 'NeighborLoader',
  'num_neighbors': [6, 2],
  'replace': True,
  'batch_size': 1024,
  'num_workers': 4
}        

produces the following plots for the performance metrics and losses.

✅ Additional Sampling methods are available at Demystifying Graph Sampling & Walk Methods

References

1. Taming PyTorch Geometric for Graph Neural Networks
2. A Practical Tutorial on Graph Neural Networks
3. A Comprehensive Introduction to Graph Neural Networks - Datacamp
4. Graph Neural Networks: A Gentil Introduction - YouTube
5. Stanford CS: Machine Learning with Graphs - YouTube
6. Github - PyTorch Geometric
7. PyTorch Geometric Datasets
8. Flickr Dataset

Patrick Nicolas has over 25 years of experience in software and data engineering, architecture design and end-to-end deployment and support with extensive knowledge in machine learning. He has been director of data engineering at Aideo Technologies since 2017 and he is the author of "Scala for Machine Learning", Packt Publishing ISBN 978-1-78712-238-3 and Geometric Learning in Python Newsletter on LinkedIn.