Source-linked AI summary
tf.data: A Machine Learning Data Processing Framework
Derek G. Murray, Jiri Simsa, Ana Klimovic, Ihor Indyk
TL;DR
Machine-learning input pipelines are challenging to implement efficiently and have received less attention despite their resource use and effect on training performance. tf.data provides a composable pipeline API and optimizing runtime; its evaluation reports substantial training-time improvements, while fleet analysis identifies diverse processing and repeated computation.
Problem
Method
tf.data offers composable operators parameterized by user functions and a runtime that graph-optimizes pipelines and tunes execution parameters automatically.
Source: §1 · p. 2, §3.3.1 · p. 6, §3.3.2 · p. 6Results
On ResNet50/ImageNet, parallelism and software pipelining reduced time to convergence by 10.4×, with caching and static optimizations improving training time by another 2×.
Source: §1 · p. 2Takeaways & Limitations
Fleet analysis across millions of jobs found diverse pipelines, substantial input-processing resource use, and repeated computation, motivating research on sharing computation and pushing projection toward storage.
Source: §1 · pp. 1–2, §7 · p. 13
Abstract
from arXiv · showhide
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.
1 Introduction
Input pipelines are difficult to optimize because they must process large, complex datasets while keeping accelerators supplied with data. tf.data combines a composable API with a runtime that optimizes pipeline execution; evaluations show both substantial training-time gains and significant fleet-wide input-processing costs.
1 Introduction
- 30% of total compute time across analyzed jobs went to ingesting data, and 20% of jobs spent more than a third of their compute time in the input pipeline.Compute time is scaled by the capability of the hardware resource.Source: §1 · p. 1, Figure 1 · p. 1
- Pipelines must ingest large volumes, apply complex transformations, overlap communication and computation, and tune parallelism and prefetching to workload and hardware.Accelerators have limited support for common preprocessing, which is commonly handled on CPUs; feeding them fast enough is increasingly challenging.Source: §1 · pp. 1–2
- tf.data provides composable, reusable operators and a runtime that represents pipelines as graphs, applies static rewrites, and tunes parallelism and prefetch buffers automatically.The abstractions let users focus on data-processing logic rather than manually tuning performance parameters.Source: §1 · p. 2
- 10.4× lower time to convergence on ResNet50/ImageNet followed parallelism and software pipelining, while caching and static optimizations improved training time by an additional 2×.The evaluation also reports that tf.data auto-tuning matches expert hand-tuned pipelines.Source: §1 · p. 2
- 13% of analyzed jobs read more than 1 TB of input data, and those jobs consumed over 96% of total compute resources.Source: Figure 2 · p. 2
2 Input Pipeline Requirements
ML input pipelines handle large datasets and transformations while meeting ordering, checkpointing, and performance needs. The section explains why flexibility and automatic optimization matter, and why existing systems do not fully meet those requirements.
Input data and transformations
- 13% of analyzed jobs read at least 1 TB of input data, so their data cannot fit in memory; jobs reading over 1 TB account for over 96% of total compute resources.Source: §2 · p. 2
- Input processing follows an extract-transform-load pattern: data is read from storage, transformed for training, and loaded onto the accelerator.Online preprocessing is the focus; offline preprocessing can extract features and convert data to formats such as Avro, Parquet, or TFRecord.Source: §2 · pp. 2–3
Data ordering
- Training pipelines must preserve ordering and batching needs, support reproducible deterministic output by seed, and allow nondeterministic execution to reduce straggler-related head-of-line blocking.They must also be checkpointable so ordering guarantees survive preemptions.Source: §2 · p. 3
Performance
- Pipelining data processing with training reduces the chance that computation waits for a batch and helps maximize accelerator utilization.Input transformations can be CPU-intensive, so throughput depends on efficiently using available CPU resources.Source: §2 · p. 3
Ease of use
- Pipeline flexibility is necessary because jobs span diverse domains and formats, with over 1,000 transformation combinations despite the 10 most common covering over 75% of jobs.The system must also optimize automatically because users may lack expertise to tune parallelism and buffer sizes.Source: §2 · p. 3
Existing systems
- Existing alternatives each miss a requirement: PyTorch DataLoader has critical-path Python and uniform-random-access assumptions, MXNet DataIter requires native extensions for new schemes, and DALI lacks workload and accelerator flexibility.Source: §2 · pp. 3–4
- General data-processing systems were not considered because TensorFlow integration could limit performance, sequential output is difficult, and extra copies can be costly for short training steps.The authors note that step times under 1 ms are not uncommon and most are under 10 ms.Source: §2 · p. 4
3 Design and Implementation
tf.data separates a statically typed dataset expression describing an input sequence from iterators that traverse it, then applies static and dynamic optimizations to execution.
3.1 Datasets and Iterators
- Datasets define typed, potentially infinite element sequences as source or transformation expression trees; iterators traverse them and support concurrent get_next calls and checkpoint save and restore.The iterator abstraction hides how elements are produced and lets optimizations change the underlying dataset graph.Source: §3.1 · p. 4
- Transformations can accept Python user-defined functions, which TensorFlow Autograph converts to dataflow graphs.Source: §3.1 · p. 4
3.2 Parallel and Distributed Execution
- Prefetch overlaps producer and consumer work, while parallel map and interleave parallelize processing and file reads; nondeterministic ordering can reduce head-of-line blocking at a reproducibility cost.Interleave can also run copies of arbitrary pipelines on different input shards to speed pipelines bottlenecked by filter or unbatch.Source: §3.2 · p. 5
- For distributed synchronous training, sharding assigns different data files to each host, and the sharded pipelines do not communicate.Source: §3.2 · p. 5
3.3 Automatic Optimization
- Static rewrites fuse adjacent transformations and vectorize map-after-batch, reducing iterator overhead and enabling batched operations.Map vectorization applies the function across batch slices in parallel.Source: §3.3.1 · pp. 5–6
- An analytical model estimates iterator output latency from processing time, input latency, buffer size, producer and consumer rates, and asynchronous parallelism.The expected latency for asynchronous transformations accounts for the probability that their buffer is empty.Source: §3.3.2 · p. 6
- The runtime periodically selects parallelism and buffer sizes to minimize end-to-end pipeline output latency under CPU and RAM budgets.The optimization targets the whole pipeline because transformations share resources and locally optimal choices can harm end-to-end performance.Source: §3.3.2 · pp. 6–7
4 Evaluation
The evaluation measures tf.data’s input-pipeline performance, its effect on end-to-end training, and its comparison with other systems. Across five MLPerf benchmarks, pipeline optimizations substantially reduce training time, while dynamic tuning approaches expert-tuned performance.
4 Evaluation
- The evaluation used MLPerf workloads spanning Resnet50, SSD, Mask-RCNN, GNMT, and Transformer, with public datasets ranging from 135 MB to 140 GB.Throughput experiments compared sequential, expert-tuned, and auto-tuned configurations; end-to-end experiments measured time to target accuracy.Source: §4 · p. 7, §4.1 · p. 8, §4.2 · p. 9
4.1 Input Pipeline Experiments
- Input-pipeline speedups ranged from 2.7× for Mask-RCNN to 63.1× for SSD against the sequential baseline.The experiments ran on a 56-core machine; parallelizing both compute and I/O can yield speedup above 56×.Source: §4.1 · p. 8, Figure 7 · p. 8
- Map-plus-batch fusion was the strongest static optimization for Resnet50 and SSD; intermediate pipeline steps prevented that fusion for Mask-RCNN and the text workloads.Non-deterministic ordering improved Resnet50 throughput by approximately 40%, while most other benchmarks improved training time by less than 4%.Source: §4.1 · p. 8
- Auto-tuned parallelism delivered performance comparable to expert hand-tuned parallelism across all tested input pipelines.For Resnet50, SSD, Mask-RCNN, and Transformer, end-to-end convergence time was within 1% of expert tuning; GNMT was within 4%.Source: §4.1 · p. 8, §4.2 · p. 9
4.2 End-to-End Experiments
- 7× average end-to-end training-time speedup across five MLPerf benchmarks came from adding input-pipeline parallelism and software pipelining over sequential input.For image workloads, further tf.data features added speedups of 2× for Resnet50, 1.2× for SSD, and 1.4× for Mask-RCNN.Source: §4.2 · p. 9, Figure 8 · p. 9
- For Resnet50, SSD, Mask-RCNN, GNMT, and Transformer, optimized tf.data pipelines matched the artificial upper bound on step performance, leaving benchmarks no longer input-bound.The upper-bound pipeline repeated a cached batch and performed no data processing.Source: §4.2 · p. 9
4.3 Comparison to Other Systems
- On a 64-core server, tf.data processed ImageNet epochs 1.9× faster than PyTorch DataLoader; DALI was faster when computation was offloaded to multiple GPUs, at added input-processing cost and GPU resource use.Disabling map-plus-batch fusion in tf.data increased processing time to 448 seconds per epoch.Source: §4.3 · p. 9, Table 6 · p. 9
- MLPerf v0.7 results show tf.data was the only input framework used across all benchmarks; the training-time comparison is not apples-to-apples because entries used different model frameworks and hardware.The reported tf.data entries trained Resnet50, SSD, Transformer, and BERT in under 30 seconds; GNMT’s higher time was attributed to slower TensorFlow model computation, not tf.data on the critical path.Source: §4.3 · p. 10, Table 7 · p. 10
5 Experience
The fleet analysis characterizes tf.data pipelines across millions of jobs, finding diverse transformations, frequent re-execution, and substantial data processing. These observations motivate research on sharing computation, processing data closer to storage, and addressing host bottlenecks and inference needs.
5 Experience
- Across millions of Google jobs studied over one month, tf.data pipelines varied widely and were frequently re-executed, motivating opportunities to reuse computation across jobs.The analysis instrumented transformations, bytes consumed and produced, iterator creations, and compute time; compute-time coverage was measured for 66% of jobs, representing 75% of total compute time, with 10% assumed for the rest.Source: §5 · p. 10, §5.1 · p. 10
5.1 Fleet-wide Input Pipeline Analysis
- Map, batch, prefetch, repeat, and zip were the most common transformations; 77% of pipelines used interleave, 87% used prefetch, and 40% used parallel map.Only 19% of jobs used the cache transformation.Source: §5.1 · pp. 10–11, Figure 9 · p. 11
- For approximately 75% of jobs, the end-to-end pipeline produced less data than it read, while map transformations expanded or reduced data by over an order of magnitude for 13% of jobs.Filter discarded more than 10% of input data volume in approximately 23% of jobs, but reduced data by less than 1% in 70% of jobs.Source: §5.1 · p. 11, Figure 10 · p. 11
- Approximately 75% of pipelines ran more than once in an hour, and 5% ran more than 100 times in an hour.The reuse analysis excluded shuffled datasets, which comprised 57% of tf.data jobs, to conservatively estimate opportunities for computation reuse.Source: §5.1 · p. 11, Figure 11 · p. 12
- Caching the top 10% of materialized datasets could capture 72% of CPU resources used to compute tf.data datasets; identical end-to-end graphs make this reuse estimate conservative.Only 10% of pipelines were re-executed across multiple jobs, and the largest observed shared pipeline ran in approximately 50,000 jobs; considering subgraphs may reveal more reuse.Source: §5.1 · pp. 11–12, Figure 12 · p. 12
5.2 Implications for Future Research
- Automated cache insertion and cross-job dataset caching are promising directions, but caching trades RAM for CPU and throughput, and randomized transformations such as shuffle should not be cached.Users must also consider effects on training quality and accelerator utilization.Source: §5.2 · p. 12
- Because preprocessing reduces data volume for 75% of jobs, applying projections closer to storage could reduce network transfers.Columnar formats can enable reading only relevant record fields.Source: §5.2 · p. 12
- When hosts cannot produce data as quickly as attached accelerators consume it, accelerators idle; proposed directions include scaling distributed input workers independently or offloading preprocessing to accelerators.Inference pipelines also need a different design: they prioritize request latency, implying less buffering, no shuffling, and different batching than training pipelines.Source: §5.2 · p. 12
6 Related Work
Related work spans data loading shared across jobs, dataflow programming and optimization, database-integrated machine learning, and automatic tuning. tf.data’s comparisons emphasize its general programming model and focus on optimizing input processing within a machine.
6 Related Work
- Prior work proposed unified data loading shared among jobs and frameworks; its observation that preprocessing can be shared agrees with tf.data’s fleet findings.tf.data instead focuses on a general programming model for diverse preprocessing schemes.Source: §6 · pp. 12–13
- tf.data draws on LINQ-family dataflow systems but adds order-sensitive transformations for ML and interoperates with existing training-loop frameworks.Unlike Optimus’s focus on network I/O in distributed queries, tf.data targets wall-clock execution of input code on a single machine; prefetch masks host-to-accelerator copies similarly to Dandelion and PTask.Source: §6 · p. 13
- Database-related systems push ML into SQL or integrate inference runtimes into query plans, whereas tf.data uses TensorFlow’s tensor representation for interoperability.Raven’s model-related optimizations are more aggressive; the passage notes they may be useful for knowledge distillation, where one model generates features for another.Source: §6 · p. 13
- Unlike SEDA’s queue-length-based thread allocation, tf.data tunes stages according to predicted end-to-end effects and models memory as a possible bottleneck.Like DS2, it uses lightweight measurements of useful processing time to make scaling decisions.Source: §6 · p. 13
7 Conclusion
The conclusion presents tf.data as a framework for composing diverse ML input pipelines and optimizing their execution. Its fleet analysis motivates future work on sharing computation and moving data projection toward storage.
7 Conclusion
- tf.data combines composable, customizable operators with dataflow-graph execution and static optimizations to improve end-to-end training time for state-of-the-art models.Parallelism and software pipelining improved Resnet50 training time by over 10×, while operator fusion provided an additional 2× improvement.Source: §7 · p. 13
- Analysis of millions of real Google jobs quantified the resource footprint, diversity, and redundant computation of ML data processing, motivating work on cross-job sharing and storage-layer projection.Source: §7 · p. 13