nnirp/NNIRP-dataset-sample

NNIRP Dataset: How You Split Is What You Get

A dataset and evaluation protocol for predicting inference runtime of neural network models from their ONNX computational graphs. Contains 103,070 profiling samples from 190 source configurations spanning 6 architecture families, organized into 156 clusters across 28 sub-families.

Dataset Summary

Each sample includes three data layers:

Architecture Families

Dataset Structure

Data is organized as one tar.gz archive per source configuration per data layer:

NNIRP-dataset/
├── manifests/
│   ├── splits.json                    # Canonical train/val/test split
│   ├── clusters.json                  # Cluster taxonomy (ID → cluster → sub-family → family)
│   └── hf_model_type_case_ids.json    # HuggingFace model_type → source config ID
├── profiling/
│   ├── 792.tar.gz                     # Profiling JSONs for source config 792
│   ├── 793.tar.gz
│   └── ...                            # 190 archives
├── pyg-features/
│   ├── 792.tar.gz                     # PyG .pt.zst files for source config 792
│   └── ...                            # 190 archives
└── onnx-graphs/
    ├── 792.tar.gz                     # ONNX .onnx files for source config 792
    └── ...                            # 190 archives

Each archive is named by its source configuration ID (integer). Source configurations are grouped into a three-level hierarchy (family → sub-family → cluster). Extracting an archive yields the sample files directly (flat, no nested directories). The manifests/clusters.json file provides the complete mapping from source configuration IDs to clusters, sub-families, and families.

Data Splits

The canonical cluster-atomic split ensures no cluster straddles two splits. All validation and test clusters satisfy a bigram coverage threshold (≥0.80) against the training pool.

Split assignments are defined in manifests/splits.json.

PyG Feature Schema

Each .pt.zst file is a zstandard-compressed PyTorch Geometric Data object:

Profiling JSON Schema

Each .json file contains summary statistics (count, mean, median, variance, min, max) for:

• —Runtime (ms) — inference latency
• —Peak VRAM (MB) — GPU memory usage
• —Peak RAM (MB) — system memory usage
• —Peak Disk Usage (MB), Disk Read (MB), Disk Write (MB)

All measurements are from an NVIDIA T4 GPU with CUDA, using PyTorch eager-mode inference.

Loading Examples

Extract and load PyG features for one source configuration

import tarfile, io, json, zstandard, torch

# Extract a single source config's PyG features
with tarfile.open("pyg-features/900.tar.gz", "r:gz") as tar:
    tar.extractall("pyg-features/900/")

# Load one sample
def load_pyg_sample(path: str):
    dctx = zstandard.ZstdDecompressor()
    with open(path, "rb") as f:
        raw = dctx.decompress(f.read())
    return torch.load(io.BytesIO(raw), weights_only=False)

data = load_pyg_sample("pyg-features/900/apple--aimv2-large-patch14-224-lit_im224_b4_fp32.pt.zst")
print(data.x.shape)           # [N, 14] node features
print(data.edge_index.shape)  # [2, E] edges
print(data.y)                 # log2(runtime_ms)

Extract and load profiling data

import tarfile, json

with tarfile.open("profiling/900.tar.gz", "r:gz") as tar:
    tar.extractall("profiling/900/")

with open("profiling/900/apple--aimv2-large-patch14-224-lit_im224_b4_fp32.json") as f:
    prof = json.load(f)
print(f"Runtime: {prof['Runtime (ms)']['mean']:.2f} ms")
print(f"VRAM: {prof['Peak VRAM (MB)']['mean']:.0f} MB")

Load split and taxonomy

import json

with open("manifests/splits.json") as f:
    splits = json.load(f)
train_ids = splits["train"]  # list of source config IDs

with open("manifests/clusters.json") as f:
    clusters = json.load(f)
# Map source config ID → family
for cluster_name, info in clusters["clusters"].items():
    family = clusters["subfamily_to_family"][info["sub_family"]]
    for config_id in info["cases"]:
        print(f"  Config {config_id}: {cluster_name} ({family})")

Download a single source configuration via the HF Hub

from huggingface_hub import hf_hub_download

# Download one archive
path = hf_hub_download(
    repo_id="nnirp/NNIRP-dataset",
    filename="pyg-features/900.tar.gz",
    repo_type="dataset",
)

Data Collection

Data was collected through a three-stage automated pipeline:

1. 1.ONNX export — neural network models are exported using a lightweight procedure that captures computational graph topology without trained weights
2. 2.GPU profiling — inference runtime, peak VRAM, and peak RAM are measured on an NVIDIA T4 GPU across multiple repetitions
3. 3.Feature encoding — ONNX graphs are converted to PyTorch Geometric Data objects with structured node, edge, and graph-level features

Parametric source configurations sweep hyperparameters (layer count, hidden dimension, batch size, precision) to generate dense scaling curves. HuggingFace source configurations profile model variants grouped by transformers model type. All models use randomly initialized weights; exported artifacts retain only graph topology and shape metadata.

Limitations

• —All profiling was performed on a single GPU type (NVIDIA T4); predictions may not generalize to other hardware without re-profiling
• —ONNX export coverage is incomplete for some operators and dynamic control flow patterns
• —Runtime measurements reflect PyTorch eager-mode inference; optimized inference engines may show different characteristics
• —Parametric source configurations account for ~22% of source configurations but ~89% of samples

License

CC-BY-NC-SA 4.0

Citation

@inproceedings{nnirp2026,
  title={How You Split Is What You Get: A Dataset and Evaluation Protocol for Neural Network Inference Runtime Prediction},
  author={Anonymous},
  booktitle={NeurIPS 2026 Evaluations and Datasets Track},
  year={2026}
}