[论文解读] tf.data: A Machine Learning Data Processing Framework

Derek G. Murray, Jiřı Šimša+2 more|arXiv (Cornell University)|Jan 28, 2021

Advanced Data Storage Technologies参考文献 48被引用 9

一句话总结

tf.data 是一个机器学习数据处理框架，通过声明式、可组合的算子和自动性能调优，实现高效、可扩展的输入管道构建。它通过应用并行处理、流水线处理、缓存和静态优化，在 ImageNet 上的 ResNet50 训练中实现了高达 10.4 倍的训练收敛速度提升，无需手动调优即可达到专家调优管道的性能水平。

ABSTRACT

Training machine learning models requires feeding input data for models to ingest. Input pipelines for machine learning jobs are often challenging to implement efficiently as they require reading large volumes of data, applying complex transformations, and transferring data to hardware accelerators while overlapping computation and communication to achieve optimal performance. We present tf.data, a framework for building and executing efficient input pipelines for machine learning jobs. The tf.data API provides operators which can be parameterized with user-defined computation, composed, and reused across different machine learning domains. These abstractions allow users to focus on the application logic of data processing, while tf.data's runtime ensures that pipelines run efficiently. We demonstrate that input pipeline performance is critical to the end-to-end training time of state-of-the-art machine learning models. tf.data delivers the high performance required, while avoiding the need for manual tuning of performance knobs. We show that tf.data features, such as parallelism, caching, static optimizations, and non-deterministic execution are essential for high performance. Finally, we characterize machine learning input pipelines for millions of jobs that ran in Google's fleet, showing that input data processing is highly diverse and consumes a significant fraction of job resources. Our analysis motivates future research directions, such as sharing computation across jobs and pushing data projection to the storage layer.

研究动机与目标

• 解决机器学习工作负载中因低效数据输入管道导致的性能瓶颈。
• 减少对性能关键参数（如并行度和预取缓冲区大小）的手动调优负担。
• 在各种机器学习工作负载（包括视觉、自然语言处理和强化学习）中实现高性能数据处理。
• 对大规模真实世界数据管道行为进行表征，以指导未来优化方向。
• 提供可重用、可组合的 API，与现有机器学习框架（如 TensorFlow 和 PyTorch）无缝集成。

提出的方法

• 通过无状态数据集和有状态迭代器暴露声明式 API，用于定义数据处理管道。
• 将管道表示为数据流图，并通过图重写应用静态优化以提高效率。
• 使用性能预测自动调优内部参数（如并行度和预取缓冲区大小）。
• 支持关键优化技术，包括软件流水线、缓存和算子融合，以重叠计算与数据加载。
• 通过将数据预处理卸载到 CPU 并重叠 I/O 与计算，与硬件加速器（如 GPU、TPU）集成。
• 使用轻量级监控机制监测处理时间和内存使用情况，支持动态扩展决策。

实验结果

研究问题

• RQ1 在现代机器学习工作负载中，输入管道性能在端到端训练时间中起到多大影响？
• RQ2 自动调优管道参数（如并行度、预取大小）在多大程度上可以达到专家调优配置的性能水平？
• RQ3 在不同工作负载和数据模态下，真实世界机器学习数据管道中的主要性能瓶颈是什么？
• RQ4 如何在大规模环境下优化数据处理工作负载，同时最小化冗余计算和资源使用？
• RQ5 在机器学习数据管道中，哪些系统级抽象和优化最有效于实现高吞吐量和低延迟？

主要发现

• 在公开基准测试中，输入管道消耗了高达 65% 的训练时间，而在 Google 集群中 30% 的总计算时间用于数据摄入。
• 仅通过并行处理和软件流水线，tf.data 即将 ResNet50 的训练收敛时间缩短 10.4 倍。
• 额外的优化（如缓存和静态图重写）使训练速度进一步提升 2 倍。
• tf.data 的自动调优实现了与专家调优管道相当的性能，且无需手动配置。
• 数百万个真实机器学习作业表现出高度多样的数据处理模式，且在相似管道间存在显著的冗余计算。
• 该分析为未来研究提供了动力，包括跨作业计算共享以及将数据投影推送到存储层，以减少预处理开销。