on the fly generation with Dataset api tensorflow
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
On-the-fly generation with tf.data means you create or transform training examples as the pipeline is consumed instead of materializing the whole dataset ahead of time. This is the standard approach when data is large, augmentation must be randomized per epoch, or examples are generated from files, records, or simulation code.
Choose the Right Dataset Source
TensorFlow gives you several starting points, and the best one depends on where the data originates.
- '
from_tensor_slicesfor arrays already in memory' - '
from_generatorfor Python generators' - file readers such as
TFRecordDatasetfor record-based storage
For genuine on-the-fly generation, from_generator is often the simplest place to start.
The important detail is output_signature. TensorFlow needs to know the shape and dtype of what the generator yields.
Use map for Cheap, Parallel Transformations
If you already have a dataset source, do not regenerate the whole example in Python just to apply a simple transformation. Use map so TensorFlow can pipeline the work efficiently.
This pattern is better than doing augmentation outside the pipeline because it keeps preprocessing close to training and allows overlap with model execution.
Batch, Shuffle, and Prefetch in the Right Order
A practical input pipeline usually looks like this:
The order matters. Shuffling before batching usually gives better sample mixing. Prefetching at the end helps overlap data preparation with training so the accelerator spends less time idle.
File-Based On-the-Fly Generation Example
For image tasks, the data often starts as file paths rather than in-memory arrays. tf.data can decode and transform each file lazily.
This avoids loading the entire image set into RAM and lets you apply transformations only when needed.
Watch the Python Boundary
from_generator is convenient, but it keeps Python in the data path. That can become a bottleneck at scale. If throughput matters, prefer pure TensorFlow ops in map, interleave, and file readers wherever possible.
A good rule is:
- start with
from_generatorwhen prototyping complex generation logic - move heavy per-example work into TensorFlow ops once correctness is established
That usually gives the best mix of iteration speed and runtime performance.
Repeat and Determinism
If training runs for multiple epochs, think deliberately about repetition and randomness.
Without repeat, the iterator is exhausted after one pass. With randomized generation, also decide whether exact reproducibility matters. If it does, seed both NumPy and TensorFlow and avoid hidden non-determinism in the generator.
Common Pitfalls
One common mistake is omitting output_signature, which leaves TensorFlow unable to build the dataset correctly. Another is doing all generation in slow Python code and then wondering why GPU utilization is poor.
People also often forget prefetch, so data preparation and training run strictly one after another. That wastes hardware.
Finally, be careful with side effects inside generators. A dataset pipeline is easier to debug when it behaves like a pure function of its inputs.
Summary
- Use
tf.datato generate or transform samples only when they are needed. - '
from_generatoris a good starting point for dynamic sample creation.' - Use
map, batching, shuffling, and prefetching to build an efficient pipeline. - Move heavy preprocessing into TensorFlow ops when Python becomes the bottleneck.
- Define shapes and dtypes explicitly with
output_signaturefor reliable pipelines.
Related reading
- On Windows, running import tensorflow generates No module named _pywrap_tensorflow error
- Online Learning with Tensorflow
- Only absolute URLs are supported when loading Keras model in Tensorflow.js with loadLayersModel
- OpenAI GPT-2 model use with TensorFlow JS
- One hot coding in Train Validation and Test set Production data
- One stage vs two stage object detection
- OnActivityResult method is deprecated, what is the alternative?
- One istio-ingressgateway and multiple TLS gateways

System Design Fundamentals
Build a strong foundation in designing scalable, reliable distributed systems.
View the courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.