on the fly generation with Dataset api tensorflow
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
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.

