Skip to content

Streaming data pipelines for graphs

Author: K3-Node Team
Backend: Multi-Backend
Dataset: GeometricShapes
Description: PyG's example builds streaming pipelines with torchdata: read raw files one by one, turn each record into a graph, and batch the graphs on the fly, without building a dataset on disk first.

View in Colab   GitHub source

Streaming data pipelines for graphs

PyG's example builds streaming pipelines with torchdata: read raw files one by one, turn each record into a graph, and batch the graphs on the fly, without building a dataset on disk first. Python generators do the same without any extra library, on every backend. This notebook builds two such pipelines: molecules from a CSV file of SMILES strings, and point clouds from mesh files.

Same model as PyG's examples/datapipe.py; with Python generators instead of torchdata, which only works with PyTorch.

Install K3-Node, then choose a backend: "tensorflow", "torch" or "jax"

!pip install k3-node[examples]
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

Molecules from a CSV file

Each row of the ESOL file holds a SMILES string and the solubility of a molecule; from_smiles turns the SMILES into a graph (with RDKit). PyG uses the larger HIV file.

import csv
import glob
import itertools

from k3_node.data import Batch
from k3_node.data.download import download_url
from k3_node.utils import from_smiles

path = download_url("https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/delaney-processed.csv", "data/raw")


def molecules(path):
    with open(path) as f:
        for row in csv.DictReader(f):
            data = from_smiles(row["smiles"])
            data.y = float(row["measured log solubility in mols per litre"])
            yield data


def batches(graphs, batch_size):
    graphs = iter(graphs)
    while chunk := list(itertools.islice(graphs, batch_size)):
        yield Batch.from_data_list(chunk)


for batch in itertools.islice(batches(molecules(path), batch_size=32), 2):
    print(batch)

Point clouds from mesh files

Every OFF mesh file is read, SamplePoints samples 1,024 points on its surface, and KNNGraph connects every point to its 6 nearest neighbors. PyG uses ModelNet10.

from k3_node.datasets import GeometricShapes
from k3_node.io import read_off
from k3_node.transforms import Compose, KNNGraph, SamplePoints

GeometricShapes("data/GeometricShapes")  # downloads the mesh files
mesh_to_graph = Compose([SamplePoints(1024), KNNGraph(k=6)])


def meshes(pattern):
    for path in sorted(glob.glob(pattern, recursive=True)):
        data = mesh_to_graph(read_off(path))
        data.category = path.split("/")[-3]
        yield data


for batch in itertools.islice(batches(meshes("data/GeometricShapes/raw/**/train/*.off"), batch_size=8), 2):
    print(batch)