TpuGraphs：大规模张量计算图的性能预测数据集

精确的硬件性能模型在代码优化中发挥着至关重要的作用。它们可以协助编译器做出启发式决策，或者帮助自动调优器识别给定程序的最佳配置。例如，XLA（一种机器学习编译器）的自动调优器在处理谷歌大量生产流量的最先进模型时，发现了10-20%的速度提升。尽管存在一些用于程序性能预测的数据集，但它们主要针对小型子程序，如基本块或内核。本文介绍了TpuGraphs，这是一个在Tensor Processing Units（TPUs）上运行的全张量程序性能预测数据集，以计算图的形式表示。数据集中每个图代表一个机器学习工作负载的主要计算部分，例如一个训练周期或一次推理步骤。每个数据样本包含一个计算图、一个编译配置以及该图在使用该配置编译后的执行时间。数据集中的图是从开源机器学习程序中收集的，涵盖了流行的模型架构，如ResNet、EfficientNet、Mask R-CNN和Transformer。TpuGraphs提供的图数量比最大的图属性预测数据集（具有相似图大小）多25倍，并且平均而言，其图规模比现有的机器学习程序性能预测数据集大770倍。这一大规模图上的图级预测任务带来了新的挑战，包括可扩展性、训练效率和模型质量等方面。