Datasets

Datasets#

MLX-graphs provides several built-in datasets for graph machine learning tasks.

Note

Datasets are automatically downloaded to .mlx_graphs_data/ in the current working directory upon first use. You can specify a custom directory by passing base_dir to the dataset constructor.

Available Datasets#

Dataset Statistics#
Dataset	Graphs	Nodes	Edges	Avg Degree	Features	Labels
Karate Club	1	34	156	9.18	node	node
Cora	1	2,708	10,556	7.80	node	node
CiteSeer	1	3,312	9,104	5.50	node	node
PubMed	1	19,717	88,651	8.99	node	node
QM7b	7,211	~15	~130	~8.67	node, edge	graph
Elliptic Bitcoin	1	203,769	234,355	2.30	node	node
MovieLens 100K	1	1,682	100,000	118.90	node	edge

Dataset Classes#

`KarateClubDataset`	Zachary's Karate Club netowork dataset from An Information Flow Model for Conflict and Fission in Small Groups.
`PlanetoidDataset`	The citation network datasets `"Cora"`, `"CiteSeer"` and `"PubMed"` from the "Revisiting Semi-Supervised Learning with Graph Embeddings" paper.
`QM7bDataset`	QM7b dataset from the "MoleculeNet: A Benchmark for Molecular Machine Learning" paper, consisting of 7,211 molecules with 14 regression targets.
`TUDataset`	A collection of over 120 benchmark datasets for graph classification and regression, made available by TU Dortmund University.
`SuperPixelDataset`	MNIST and CIFAR10 superpixel datasets for graph classification tasks converted fromt the original MINST and CIFAR10 images.
`OGBDataset`	Datasets from the Open Graph Benchmark (OGB) collection of realistic, large-scale, and diverse benchmark datasets for machine learning on graphs.
`EllipticBitcoinDataset`	The Elliptic Bitcoin dataset of Bitcoin transactions from the "Anti-Money Laundering in Bitcoin: Experimenting with Graph Convolutional Networks for Financial Forensics" paper.
`MovieLens100K`	The MovieLens 100K heterogeneous rating dataset, assembled by GroupLens Research from the MovieLens web site, consisting of movies (1,682 nodes) and users (943 nodes) with 100K ratings between them.
`IMDB`	A subset of the Internet Movie Database (IMDB), as collected in the "MAGNN: Metapath Aggregated Graph Neural Network for Heterogeneous Graph Embedding" paper.
`DBLP`	A subset of the DBLP computer science bibliography website, as collected in the "MAGNN: Metapath Aggregated Graph Neural Network for Heterogeneous Graph Embedding" paper.

KarateClubDataset#

class mlx_graphs.datasets.KarateClubDataset[source]#

Bases: Dataset

Zachary’s Karate Club netowork dataset from An Information Flow Model for Conflict and Fission in Small Groups. This is a simple dataset for node classification. The graph has 34 nodes and 156 (undirected) edges. Each node belongs to one of 2 classes.

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

PlanetoidDataset#

class mlx_graphs.datasets.PlanetoidDataset(name: Literal['cora', 'citeseer', 'pubmed'], split: Literal['public', 'full', 'geom-gcn'] = 'public', without_self_loops: bool = True, base_dir: str | None = None)[source]#

Bases: Dataset

The citation network datasets "Cora", "CiteSeer" and "PubMed" from the “Revisiting Semi-Supervised Learning with Graph Embeddings” paper. Nodes represent documents and edges represent citation links. Training, validation and test splits are given by binary masks.

This dataset follows a similar implementation as in PyG.

Parameters:

name (Literal['cora', 'citeseer', 'pubmed']) – The name of the dataset ("Cora", "CiteSeer", "PubMed").
split (str, optional) –
The type of dataset split ("public", "full", "geom-gcn"). If set to "public", the split will be the public fixed split from the “Revisiting Semi-Supervised Learning with Graph Embeddings” paper. If set to "full", all nodes except those in the validation and test sets will be used for training (as in the “FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling” paper). If set to "geom-gcn", the 10 public fixed splits from the “Geom-GCN: Geometric Graph Convolutional Networks” paper are given.
without_self_loops (bool) – Whether to remove self loops. Default to True.
base_dir (Optional[str]) – Directory where to store dataset files. Default is in the local directory .mlx_graphs_data/.

Example:

from mlx_graphs.datasets import Planetoid

dataset = Planetoid("cora")
>>> cora(num_graphs=1)

dataset[0]
>>> GraphData(
        edge_index(shape=(2, 10556), int32)
        node_features(shape=(2708, 1433), float32)
        node_labels(shape=(2708,), int32)
        train_mask(shape=(2708,), bool)
        val_mask(shape=(2708,), bool)
        test_mask(shape=(2708,), bool))

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property processed_path: str#: The path where raw files are stored. Defaults at <base_dir>/<name>/processed

property raw_path: str#: The path where raw files are stored.

QM7bDataset#

class mlx_graphs.datasets.QM7bDataset(base_dir: str | None = None)[source]#

Bases: Dataset

QM7b dataset from the “MoleculeNet: A Benchmark for Molecular Machine Learning” paper, consisting of 7,211 molecules with 14 regression targets.

Parameters:: base_dir (Optional[str]) – Directory where to store dataset files. Default is in the local directory .mlx_graphs_data/.

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

TUDataset#

class mlx_graphs.datasets.TUDataset(name: str, cleaned: bool = False, base_dir: str | None = None)[source]#

Bases: Dataset

A collection of over 120 benchmark datasets for graph classification and regression, made available by TU Dortmund University. Access all these datasets here.

This class also supports cleaned dataset versions containing only non-isomorphic graphs, and presented in Understanding Isomorphism Bias in Graph Data Sets.

Parameters:

name (str) – Name of the dataset to load (e.g. “MUTAG”, “PROTEINS”, “IMDB-BINARY”, etc.).
cleaned (bool) – Whether to use the cleaned or original version of datasets. Default is False.
base_dir (Optional[str]) – Directory where to store dataset files. Default is in the local directory .mlx_graphs_data/.

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property processed_path: str#: The path where raw files are stored. Defaults at <base_dir>/<name>/processed

property raw_path: str#: The path where raw files are stored.

SuperPixelDataset#

class mlx_graphs.datasets.SuperPixelDataset(name: Literal['MNIST', 'CIFAR10'], split: Literal['train', 'test'], use_features: bool = False, base_dir: str | None = None)[source]#

Bases: Dataset

MNIST and CIFAR10 superpixel datasets for graph classification tasks converted fromt the original MINST and CIFAR10 images.

The datasets were introduced in http://arxiv.org/abs/2003.00982.

Parameters:

name (Literal['MNIST', 'CIFAR10']) – name of the selected dataset
split (Literal['train', 'test']) – split of the dataset to load
use_features (bool) – if True, the adjacency matrix is computed from superpixels locations and features. If False, only from superpixels locations. Defaults to False.
base_dir (Optional[str]) – directory where to store the datasets

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property processed_path: str#: The path where raw files are stored. Defaults at <base_dir>/<name>/processed

OGBDataset#

class mlx_graphs.datasets.OGBDataset(name: Literal['ogbn-products', 'ogbn-proteins', 'ogbn-arxiv', 'ogbn-papers100M', 'ogbl-ppa', 'ogbl-collab', 'ogbl-ddi', 'ogbl-citation2', 'ogbl-vessel', 'ogbg-molhiv', 'ogbg-molpcba', 'ogbg-ppa', 'ogbg-code2'], split: Literal['train', 'val', 'test'] | None = None, base_dir: str | None = None)[source]#

Bases: Dataset

Datasets from the Open Graph Benchmark (OGB) collection of realistic, large-scale, and diverse benchmark datasets for machine learning on graphs.

Datasets belongs to three fundamental graph machine learning task categories: predicting the properties of nodes, links, and graphs. Node property prediction datasets consist of a single graph with three additional properties: train_mask, val_mask and test_mask specifying the masks for the train, validation and test splits. Link property prediction datasets also consist of a single graphs with three additional properties: train_edge_index, val_edge_index and test_edge_index, specifying the edges to be considered for training, validation and testing. Graph property prediction datasets consists of multiple graphs. The desired split can be specified via the split arg.

See here for further details and a list of the available datasets with their descriptions

Parameters:

name (Literal['ogbn-products', 'ogbn-proteins', 'ogbn-arxiv', 'ogbn-papers100M', 'ogbl-ppa', 'ogbl-collab', 'ogbl-ddi', 'ogbl-citation2', 'ogbl-vessel', 'ogbg-molhiv', 'ogbg-molpcba', 'ogbg-ppa', 'ogbg-code2']) – Name of the dataset
split (Optional[Literal['train', 'val', 'test']]) – Split of the dataset to load. Thi parameter has effect only on graph property prediction dataset. If None, the entire dataset is loaded. Defaults to None.
base_dir (Optional[str]) – Directory where to store dataset files. Default is in the local directory .mlx_graphs_data/.

Note

ogb needs to be installed to use this dataset

Note

The ogbn-mag, ogbl-wikikg2 and igbl-biokg and the graphs belonging to the largs-scale challenge category are currently not available as they require heterogenous graphs which are not yet supported by mlx-graphs

Example:

from mlx_graphs.datasets.ogb_dataset import OGBDataset

ds = OGBDataset("ogbg-molhiv", split="train")
>>> ogbg-molhiv(num_graphs=32901)

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property processed_path: str#: The path where raw files are stored. Defaults at <base_dir>/<name>/processed

EllipticBitcoinDataset#

class mlx_graphs.datasets.EllipticBitcoinDataset(base_dir: str | None = None, pre_transform: Callable | None = None, transform: Callable | None = None)[source]#

Bases: Dataset

The Elliptic Bitcoin dataset of Bitcoin transactions from the “Anti-Money Laundering in Bitcoin: Experimenting with Graph Convolutional Networks for Financial Forensics” paper.

EllipticBitcoinDataset maps Bitcoin transactions to real entities belonging to licit categories (exchanges, wallet providers, miners, licit services, etc.) versus illicit ones (scams, malware, terrorist organizations, ransomware, Ponzi schemes, etc.)

There exists 203,769 node transactions and 234,355 directed edge payments flows, with two percent of nodes (4,545) labelled as illicit, and twenty-one percent of nodes (42,019) labelled as licit. The remaining transactions are unknown

Parameters:

base_dir (Optional[str]) – Directory where to store dataset files. Default is in the local directory .mlx_graphs_data/.
pre_transform (Optional[Callable]) – A function/transform which takes in a GraphData object and returns a transformed version. The data will be transformed before saving to the disk.
transforms – A function/transform that takes in a graphData object and returns a transformed version The data object will be transformed before every access

download()[source]#: Download the dataset at self.raw_path.

process(train=True)[source]#: Process the dataset and store data in self.data

MovieLens100K#

class mlx_graphs.datasets.MovieLens100K(base_dir: str, transform: Callable | None = None, pre_transform: Callable | None = None)[source]#

Bases: HeteroDataset

The MovieLens 100K heterogeneous rating dataset, assembled by GroupLens Research from the MovieLens web site, consisting of movies (1,682 nodes) and users (943 nodes) with 100K ratings between them. User ratings for movies are available as ground truth labels. Features of users and movies are encoded according to the “Inductive Matrix Completion Based on Graph Neural Networks” paper.

Parameters:

base_dir (str) – Directory where to store dataset files.
transform (callable, optional) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before every access. (default: None)
pre_transform (callable, optional) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property raw_path: str#: The path where raw files are stored.

IMDB#

class mlx_graphs.datasets.IMDB(base_dir: str | None = None, transform: Callable | None = None, pre_transform: Callable | None = None)[source]#

Bases: HeteroDataset

A subset of the Internet Movie Database (IMDB), as collected in the “MAGNN: Metapath Aggregated Graph Neural Network for Heterogeneous Graph Embedding” paper. IMDB is a heterogeneous graph containing three types of entities - movies (4,278 nodes), actors (5,257 nodes), and directors (2,081 nodes). The movies are divided into three classes (action, comedy, drama) according to their genre. Movie features correspond to elements of a bag-of-words representation of its plot keywords.

Parameters:

base_dir (Optional[str]) – directory where the dataset should be saved.
transform (Optional[Callable]) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before every access. (default: None)
pre_transform (Optional[Callable]) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property raw_path: str#: The path where raw files are stored.

DBLP#

class mlx_graphs.datasets.DBLP(base_dir: str | None = None, transform: Callable | None = None, pre_transform: Callable | None = None)[source]#

Bases: HeteroDataset

A subset of the DBLP computer science bibliography website, as collected in the “MAGNN: Metapath Aggregated Graph Neural Network for Heterogeneous Graph Embedding” paper. DBLP is a heterogeneous graph containing four types of entities - authors (4,057 nodes), papers (14,328 nodes), terms (7,723 nodes), and conferences (20 nodes). The authors are divided into four research areas (database, data mining, artificial intelligence, information retrieval). Each author is described by a bag-of-words representation of their paper keywords.

Parameters:

base_dir (Optional[str]) – directory where the dataset should be saved.
transform (Optional[Callable]) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before every access. (default: None)
pre_transform (Optional[Callable]) – A function/transform that takes in an HeteroGraphData object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

download()[source]#: Download the dataset at self.raw_path.

process()[source]#: Process the dataset and store data in self.data

property raw_path: str#: The path where raw files are stored.

Using Datasets#

Loading a dataset is straightforward:

from mlx_graphs.datasets import KarateClubDataset, PlanetoidDataset

# Simple built-in dataset
dataset = KarateClubDataset()
graph = dataset[0]
print(f"Nodes: {graph.num_nodes}, Edges: {graph.num_edges}")

# Citation network dataset
cora = PlanetoidDataset(name="cora")
print(cora[0])

Computing Dataset Statistics#

To compute statistics for any dataset, use the included utility script:

# Stats for all available datasets
python examples/compute_dataset_stats.py

# Stats for a specific dataset
python examples/compute_dataset_stats.py --dataset karate_club

# Output as markdown table (for use in docs/issues)
python examples/compute_dataset_stats.py --format markdown

# Output as RST table (for Sphinx docs)
python examples/compute_dataset_stats.py --format rst

Datasets

Contents

Datasets#

Available Datasets#

Dataset Classes#

KarateClubDataset#

PlanetoidDataset#

QM7bDataset#

TUDataset#

SuperPixelDataset#

OGBDataset#

EllipticBitcoinDataset#

MovieLens100K#

IMDB#

DBLP#

Using Datasets#

Computing Dataset Statistics#