TensorFlow
Dataset API
on-the-fly data
machine learning
data processing

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_slices for arrays already in memory'
  • 'from_generator for Python generators'
  • file readers such as TFRecordDataset for record-based storage

For genuine on-the-fly generation, from_generator is often the simplest place to start.

python
1import numpy as np
2import tensorflow as tf
3
4
5def sample_generator():
6    for _ in range(1000):
7        x = np.random.rand(8).astype("float32")
8        y = np.array(np.sum(x) > 4.0, dtype="int32")
9        yield x, y
10
11
12dataset = tf.data.Dataset.from_generator(
13    sample_generator,
14    output_signature=(
15        tf.TensorSpec(shape=(8,), dtype=tf.float32),
16        tf.TensorSpec(shape=(), dtype=tf.int32),
17    ),
18)
19
20for features, label in dataset.take(3):
21    print(features.shape, label.numpy())

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.

python
1def add_noise(x, y):
2    noise = tf.random.normal(tf.shape(x), stddev=0.05)
3    return x + noise, y
4
5
6dataset = dataset.map(add_noise, num_parallel_calls=tf.data.AUTOTUNE)

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:

python
dataset = dataset.shuffle(1000)
dataset = dataset.batch(32)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

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.

python
1files = tf.data.Dataset.list_files("images/*.jpg", shuffle=True)
2
3
4def load_image(path):
5    bytes_ = tf.io.read_file(path)
6    image = tf.image.decode_jpeg(bytes_, channels=3)
7    image = tf.image.resize(image, [128, 128])
8    image = tf.cast(image, tf.float32) / 255.0
9    label = tf.constant(0, dtype=tf.int32)
10    return image, label
11
12
13dataset = files.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
14dataset = dataset.batch(16).prefetch(tf.data.AUTOTUNE)

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_generator when 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.

python
dataset = dataset.repeat()

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.data to generate or transform samples only when they are needed.
  • 'from_generator is 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_signature for reliable pipelines.

Course illustration
Course illustration

All Rights Reserved.