# tf.data: A Machine Learning Data Processing Framework

arXiv: https://arxiv.org/abs/2101.12127 · HTML: https://mlsanity.com/paper/2101.12127/

## Key facts

- Cited by 41 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 64 references that are in the map.
- First posted on arXiv 28 Jan 2021.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [SSD: Single Shot MultiBox Detector](https://mlsanity.com/paper/1512.02325/) (2015)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [MLPerf Training Benchmark](https://mlsanity.com/paper/1910.01500/) (2019)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [Analyzing and Mitigating Data Stalls in DNN Training](https://mlsanity.com/paper/2007.06775/) (2020)
- [Faster Neural Network Training with Data Echoing](https://mlsanity.com/paper/1907.05550/) (2019)
- [Extending Relational Query Processing with ML Inference](https://mlsanity.com/paper/1911.00231/) (2019)
- [Jointly Optimizing Preprocessing and Inference for DNN-based Visual Analytics](https://mlsanity.com/paper/2007.13005/) (2020)
- [AutoGraph: Imperative-style Coding with Graph-based Performance](https://mlsanity.com/paper/1810.08061/) (2018)
- [Helix: Accelerating Human-in-the-loop Machine Learning](https://mlsanity.com/paper/1808.01095/) (2018)

## Built on by (most cited first, 30 of 41)

- [DataPerf: Benchmarks for Data-Centric AI Development](https://mlsanity.com/paper/2207.10062/) (2022)
- [FFCV: Accelerating Training by Removing Data Bottlenecks](https://mlsanity.com/paper/2306.12517/) (2023)
- [Characterization of Large Language Model Development in the Datacenter](https://mlsanity.com/paper/2403.07648/) (2024)
- [Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learning](https://mlsanity.com/paper/2307.09619/) (2023)
- [Understanding Data Storage and Ingestion for Large-Scale Deep Recommendation Model Training](https://mlsanity.com/paper/2108.09373/) (2021)
- [Plumber: Diagnosing and Removing Performance Bottlenecks in Machine Learning Data Pipelines](https://mlsanity.com/paper/2111.04131/) (2021)
- [End-to-end Optimization of Machine Learning Prediction Queries](https://mlsanity.com/paper/2206.00136/) (2022)
- [tf.data service: A Case for Disaggregating ML Input Data Processing](https://mlsanity.com/paper/2210.14826/) (2022)
- [Where Is My Training Bottleneck? Hidden Trade-Offs in Deep Learning Preprocessing Pipelines](https://mlsanity.com/paper/2202.08679/) (2022)
- [cedar: Optimized and Unified Machine Learning Input Data Pipelines](https://mlsanity.com/paper/2401.08895/) (2024)
- [Modyn: Data-Centric Machine Learning Pipeline Orchestration](https://mlsanity.com/paper/2312.06254/) (2023)
- [MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training](https://mlsanity.com/paper/2407.12117/) (2024)
- [TensorSocket: Shared Data Loading for Deep Learning Training](https://mlsanity.com/paper/2409.18749/) (2024)
- [MegaScale-Data: Scaling Dataloader for Multisource Large Foundation Model Training](https://mlsanity.com/paper/2504.09844/) (2025)
- [InTune: Reinforcement Learning-based Data Pipeline Optimization for Deep Recommendation Models](https://mlsanity.com/paper/2308.08500/) (2023)
- [I/O in Machine Learning Applications on HPC Systems: A 360-degree Survey](https://mlsanity.com/paper/2404.10386/) (2024)
- [The Streaming Batch Model for Efficient and Fault-Tolerant Heterogeneous Execution](https://mlsanity.com/paper/2501.12407/) (2025)
- [Mixtera: A Data Plane for Foundation Model Training](https://mlsanity.com/paper/2502.19790/) (2025)
- [ShadowServe: Interference-Free KV Cache Fetching for Distributed Prefix Caching](https://mlsanity.com/paper/2509.16857/) (2025)
- [PreSto: An In-Storage Data Preprocessing System for Training Recommendation Models](https://mlsanity.com/paper/2406.14571/) (2024)
- [Machine Learning Fleet Efficiency: Analyzing and Optimizing Large-Scale Google TPU Systems with ML Productivity Goodput](https://mlsanity.com/paper/2502.06982/) (2025)
- [MinatoLoader: Accelerating Machine Learning Training Through Efficient Data Preprocessing](https://mlsanity.com/paper/2509.10712/) (2025)
- [Bullion: A Column Store for Machine Learning](https://mlsanity.com/paper/2404.08901/) (2024)
- [GPUs, CPUs, and... NICs: Rethinking the Network's Role in Serving Complex AI Pipelines](https://mlsanity.com/paper/2502.15712/) (2025)
- [Analytics Modelling over Multiple Datasets using Vector Embeddings](https://mlsanity.com/paper/2502.17060/) (2025)
- [Synergy: Resource Sensitive DNN Scheduling in Multi-Tenant Clusters](https://mlsanity.com/paper/2110.06073/) (2021)
- [A Machine Learning Data Fusion Model for Soil Moisture Retrieval](https://mlsanity.com/paper/2206.09649/) (2022)
- [L3: Accelerator-Friendly Lossless Image Format for High-Resolution, High-Throughput DNN Training](https://mlsanity.com/paper/2208.08711/) (2022)
- [Accelerating Transfer Learning with Near-Data Computation on Cloud Object Stores](https://mlsanity.com/paper/2210.08650/) (2022)
- [KerasCV and KerasNLP: Vision and Language Power-Ups](https://mlsanity.com/paper/2405.20247/) (2024)

## Abstract (arXiv)

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.