nnirp/NNIRP-dataset

NNIRP数据集是一个用于从ONNX计算图中预测神经网络模型推理运行时的数据集和评估协议。它包含103,070个分析样本，来自6个架构家族的190个源配置，组织成28个子家族的156个集群。每个样本包括三个数据层：分析数据（.json，包含运行时、VRAM和RAM统计信息）、PyG特征（.pt.zst，包含节点、边和图级特征的PyTorch Geometric图编码）和ONNX图（.onnx，仅包含拓扑结构的轻量级ONNX计算图）。数据集结构按源配置和数据层组织，包含一个tar.gz存档。数据分割采用集群原子分割，确保没有集群跨越两个分割。数据集还提供了详细的PyG特征和分析JSON的模式描述，以及如何加载和使用数据集的示例。数据收集通过一个三阶段的自动化流程完成，包括ONNX导出、GPU分析和特征编码。数据集的主要局限性包括所有分析均在NVIDIA T4 GPU上进行，ONNX导出覆盖不完整，以及运行时测量反映PyTorch eager-mode推理。数据集采用CC-BY-NC-SA 4.0许可证。

The NNIRP dataset is a dataset and evaluation protocol for predicting inference runtime of neural network models from their ONNX computational graphs. It contains 103,070 profiling samples from 190 source configurations spanning 6 architecture families, organized into 156 clusters across 28 sub-families. Each sample includes three data layers: profiling data (.json, containing runtime, VRAM, and RAM statistics), PyG features (.pt.zst, containing PyTorch Geometric graph encodings with node, edge, and graph-level features), and ONNX graphs (.onnx, lightweight ONNX computational graphs with topology only). The dataset is organized as one tar.gz archive per source configuration per data layer. Data splits are cluster-atomic, ensuring no cluster straddles two splits. The dataset also provides detailed schemas for PyG features and profiling JSONs, along with examples of how to load and use the dataset. Data was collected through a three-stage automated pipeline: ONNX export, GPU profiling, and feature encoding. Key limitations include profiling performed on a single GPU type (NVIDIA T4), incomplete ONNX export coverage, and runtime measurements reflecting PyTorch eager-mode inference. The dataset is licensed under CC-BY-NC-SA 4.0.

搜集汇总

数据集介绍

构建方式

NNIRP-dataset通过三阶段自动化流水线构建而成：首先从190个源配置中提取神经网络模型的ONNX计算图，覆盖注意力、卷积、循环等6大架构家族及28个子家族，共生成103,070个样本；随后在NVIDIA T4 GPU上测量推理延迟、显存和内存峰值等运行时指标；最后将ONNX图转换为PyTorch Geometric格式的图结构数据，包含节点、边和图级特征。所有模型仅保留拓扑结构与形状元数据，未存储训练权重，从而在保证信息完整性的同时大幅压缩存储开销。

特点

该数据集的核心特点在于其分层组织与原子化划分策略。样本按源配置ID归档，并通过家族-子家族-集群三级分类体系进行结构化管理。官方提供的规范划分保证每个集群完整归属单一子集，同时确保验证集与测试集对训练集的二元覆盖阈值不低于0.80，有效避免了分布偏移。每个样本集成了ONNX图、PyG特征向量与详细性能基准数据三层信息，为图神经网络驱动的运行时预测提供了多模态基础。

使用方法

研究者可通过HuggingFace Hub下载所需源配置的压缩归档文件（tar.gz），并利用提供的Python示例代码分别提取性能日志、PyG特征向量或ONNX图数据。通过加载manifests/splits.json与clusters.json文件可获取标准训练/验证/测试划分及其分类归属信息。PyG特征以zstandard压缩的PyTorch张量文件存储，可直接加载为Data对象用于模型训练。性能日志则包含运行时和内存指标的统计量（均值、中位数等），便于进行回归任务分析。

背景与挑战

背景概述

随着深度学习模型在规模和复杂性上的持续增长，其在实际部署中的计算开销——尤其是推理延迟和显存占用——已成为制约应用落地的关键瓶颈。在此背景下，NNIRP-dataset于2026年由匿名研究团队在NeurIPS 2026 Evaluations and Datasets Track上提出，旨在为神经网络推理运行时预测任务提供标准化评测基准。该数据集从神经网络模型的ONNX计算图出发，包含来自6个架构家族（包括注意力编码器、解码器、卷积、循环网络等）、28个子类、156个聚类的190个源配置，总计103,070个样本。每个样本涵盖用户级性能剖析数据、图结构特征（基于PyTorch Geometric）以及轻量级ONNX计算图拓扑信息，为从计算图到硬件表现之间的映射关系研究提供了系统化资源，对模型效率估计、自动架构搜索和资源感知部署等方向具有显著推动作用。

当前挑战

该数据集主要面临三方面挑战。首先，在领域问题层面，神经网络推理运行时预测的核心困难在于计算图拓扑、操作符类型、数据流特征与硬件执行行为之间复杂的非线性关系，传统基于手工规则或简单统计特征的方法难以捕捉其内在规律；该数据集通过提供结构化的图特征和细粒度性能标签，为发展数据驱动的预测模型奠定基础。其次，在构建过程中，研究者需在无训练权重的条件下忠实地表征计算图拓扑，同时处理ONNX导出中对某些动态控制流操作符覆盖不完整的问题；此外，跨不同源配置（如层数、隐藏维度、批大小、精度等参数化搜索空间）的性能均在同一单GPU（NVIDIA T4）上采集，限制了模型对异构硬件的泛化能力，而参数化配置虽仅占源配置的22%却提供了约89%的样本，导致样本分布存在偏斜，缓解这一分布失衡将成为后续研究中的重要挑战。

常用场景

经典使用场景

在深度学习系统优化与推理部署领域，NNIRP-dataset为预测神经网络模型在特定硬件上的推理运行时性能提供了规范化的数据基础。研究者可利用该数据集训练基于图神经网络的回归模型，输入模型的计算图拓扑结构，输出推断延迟与峰值显存占用等关键性能指标，从而实现对未见过架构的运行时行为的精准预估。

实际应用

在实际工业部署中，NNIRP-dataset可用于硬件选型评估、推理引擎配置优化以及资源预算规划。凭借对模型计算图的ONNX格式抽象与运行时性能标注，开发人员能在不实际运行模型的情况下预判其在不同GPU或批处理设置下的性能表现，显著压缩模型上线的调试周期。此外，该数据集也为异构计算集群中的模型-资源映射决策提供了可量化的性能预测依据。

衍生相关工作

围绕NNIRP-dataset，学界已涌现出一系列基础性工作，其中包括基于图注意力网络的推理时间预测框架、面向计算图拓扑特征的性能建模专用图神经网络架构，以及针对跨架构泛化性能的系统性基准评估协议。这些衍生研究不仅验证了该数据集在多架构性能预测中的有效性，还推动了诸如模型-硬件协同设计与自适应推理引擎选择等后续研究方向的发展。