Introduction

Preparing graph datasets for machine learning often involves tedious, non-standardized processes that hinder reproducibility and benchmarking. The OGB DatasetSaver, part of the Open Graph Benchmark (OGB) ecosystem, solves this by providing a unified interface for creating datasets compatible with the OGB framework. With the OGB library gaining massive traction in the graph ML community, this tool ensures that new datasets can be seamlessly integrated into standardized pipelines for node, link, and graph property prediction.

What Is OGB DatasetSaver?

OGB DatasetSaver is a Python class within the ogb.io module designed to help external contributors and researchers prepare their graph datasets in an OGB-compatible format. It acts as a serialization layer that ensures graph structures, node/edge features, and dataset splits are saved in a way that OGB’s data loaders can read them without further modification.

The tool is maintained as part of the snap-stanford/ogb repository, licensed under the Apache License 2.0, and is primarily written in Python. It is specifically engineered to support the three fundamental graph ML tasks: node property prediction, node property prediction, and graph property prediction.

Why OGB DatasetSaver Matters

Before the standardization provided by OGB, graph ML researchers often shared datasets in disparate formats (CSV, JSON, custom binary files), making it difficult to compare models across different benchmarks. The DatasetSaver class fills this gap by enforcing a strict schema for how graph data is stored, which allows the community to use unified evaluators and leaderboards.

By using DatasetSaver, developers can ensure their datasets are “plug-and-play” with popular frameworks like PyTorch Geometric (PyG) and the Deep Graph Library (DGL). This reduces the time spent on data engineering and allows researchers to focus on model architecture and algorithmic innovation rather than parsing custom file formats.

Key Features

• Standardized Serialization: Ensures that graph data follows the OGB convention, making it instantly compatible with OGB data loaders.
• Support for Heterogeneous Graphs: Provides specific dictionary-based mapping for heterogeneous graphs, allowing for multiple node and edge types.
• Unified Split Management: Simplifies the creation of train, validation, and test splits, ensuring that the data leakage is avoided through standardized indexing.
• Automated Metadata Generation: Automatically handles the creation of metadata files that describe the dataset version, task type, and graph properties.
• Framework Agnostic: While it integrates perfectly with PyG and DGL, the saved files are library-agnostic, meaning they can be used with TensorFlow or MxNet.
• Flexible Feature Mapping: Supports optional node and edge features, allowing researchers to define the dimensionality of their feature vectors independently.

How OGB DatasetSaver Compares

When preparing graph datasets, researchers typically choose between custom scripts, general-purpose graph libraries, or specialized benchmark tools like OGB. DatasetSaver is unique because it focuses on benchmark compatibility rather than just data storage.

The primary differentiator is that OGB DatasetSaver is not just a saving mechanism; it is a protocol. While NetworkX is excellent for graph manipulation, it does not enforce the specific splits and metadata required for a competitive benchmark. Custom scripts often lead to “hidden” data leakage or inconsistent evaluation metrics. DatasetSaver eliminates these risks by forcing the user to follow the OGB convention (e.g., naming datasets starting with ogbn-, ogbl-, or ogbg-).

Getting Started: Installation

To use the DatasetSaver class, you must install the ogb Python package. The library is compatible with Python 3.6+ and requires several dependencies including NumPy and Pandas.

Pip Installation

The recommended way to install OGB is via pip:

pip install ogb

To ensure you are using the latest version (1.3.6 or newer), you can run:

pip install -U ogb

From Source

If you intend to contribute to the OGB project or modify the internal saving logic, install from the GitHub source:

git clone https://github.com/snap-stanford/ogb
cd ogb
pip install -e .

Prerequisites

Ensure you have the following installed: Python >= 3.6, NumPy >= 1.16.0, and Pandas >= 0.24.0. If you are using the datasets with PyG or DGL, ensure those frameworks are installed in your environment.

How to Use OGB DatasetSaver

Using DatasetSaver follows a strict linear workflow. You must execute the saving methods in the exact order specified by the OGB documentation to ensure the final dataset files are generated correctly.

The process begins by initializing the DatasetSaver object with a dataset name that follows the OGB naming convention (e.g., ogbg- for graph property prediction). You then save the graph list, followed by the target labels, and finally the dataset splits.

For homogeneous graphs, you provide a list of dictionaries where each dictionary contains the edge_index and num_nodes. For heterogeneous graphs, you use dictionaries mapping node and edge types to their respective indices and features.

Code Examples

The following examples demonstrate how to prepare a toy graph property prediction dataset using the OGB DatasetSaver.

Basic Homogeneous Graph Dataset

from ogb.io import DatasetSaver
import numpy as np
import networkx as nx

# 1. Initialize the saver
dataset_name = 'ogbg-toy'
saver = DatasetSaver(dataset_name=dataset_name, is_hetero=False, version=1)

# 2. Create a list of graph objects
graph_list = []
for i in range(100):
    g = nx.fast_gnp_random_graph(10, 0.5)
    graph = {
        'edge_index': np.array(g.edges).transpose(),
        'num_nodes': len(g.nodes),
        'node_feat': np.random.randn(len(g.nodes), 3),
        'edge_feat': np.random.randn(len(g.edges), 3)
    }
    graph_list.append(graph)

# 3. Save the graphs
saver.save_graph_list(graph_list)

# 4. Save the labels
labels = np.random.randint(0, 2, size=(100, 1))
saver.save_target_labels(labels)

# 5. Save the splits
split_idx = {
    'train': np.arange(0, 80),
    'valid': np.arange(80, 90),
    'test': np.arange(90, 100)
}
saver.save_split(split_idx, split_name='random')

In this example, we use NetworkX to generate random graphs and then convert them into the dictionary format required by DatasetSaver. The edge_index must be a NumPy array of shape (2, num_edges) and num_nodes must be an integer.

Real-World Use Cases

• Molecular Property Prediction: A chemist using RDKit to process molecular graphs can use DatasetSaver to save their data in a format that any GNN researcher can load using PygGraphPropPredDataset.
• Social Network Analysis: A researcher analyzing a new social network can standardize their node classification task (ogbn-) to ensure their results are comparable to the ogbn-arxiv or ogbn-products benchmarks.
• Knowledge Graph Completion: For heterogeneous data, a developer can use the is_hetero=True flag to save complex relational data that is compatible with OGB’s heterogeneous data loaders.

Contributing to OGB DatasetSaver

The OGB project is an ongoing effort by Stanford SNAP and other collaborators. Contributions are welcome through the standard GitHub flow. Users can report bugs via the GitHub Issues tab or suggest new benchmark datasets through the project’s discussions.

If you are submitting a new dataset to be included in the official OGB library, ensure you use the DatasetSaver class to generate the files, as this is the required protocol for official inclusion.

Community and Support

The OGB community is centered around the official GitHub repository and the OGB website. For technical support, the GitHub Discussions forum is the primary channel for collaborating with the developer community and asking questions about data preparation.

Documentation is available on the official OGB website and within the ogb/io/README.md file in the repository. The project is widely used by the graph ML community, with thousands of stars and a high level of activity in its issues and discussions.

Conclusion

The OGB DatasetSaver class is a critical piece of infrastructure for the graph machine learning community. By enforcing a standardized format for dataset creation, it eliminates the friction of data engineering and enables a fair, reproducible comparison of GNN models. Whether you are creating a new benchmark or simply organizing your internal graph data, DatasetSaver provides the necessary tools to ensure your data is compatible with the state-of-the-art loaders in PyG and DGL.

To get started, install the ogb package, follow the quickstart guide in the ogb.io module, and try creating your own toy dataset. Star the repo, join the community, and contribute your datasets to the rest of the world.

Resources

• Official OGB GitHub Repository
• Open Graph Benchmark Official Website
• DatasetSaver Documentation