TensorFlow
training data
on-the-fly data generation
machine learning
data augmentation

How can I generate training data on the fly in 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.

Practice ML system design

Introduction

Generating training data on the fly in TensorFlow is a strong pattern for large datasets, procedural augmentation, and experiments where storing every transformed sample is expensive. In practice, the fastest path is to reduce the problem to a small reproducible baseline first, then reintroduce production constraints one by one. That approach keeps debugging local, prevents overfitting to one failing symptom, and makes your final implementation easier to explain to teammates.

The core design choice is where randomness and preprocessing happen. Keep expensive CPU transforms in the input pipeline, keep deterministic labels aligned with each sample, and avoid Python-side bottlenecks that starve the GPU. A strong implementation separates configuration from execution flow, adds measurable checkpoints, and captures enough telemetry to distinguish transient failures from deterministic misconfiguration.

Core Sections

1) Define a narrow baseline before optimization

Start by identifying the smallest end-to-end version that should work reliably. Keep external dependencies minimal, remove optional features, and make defaults explicit. Once the baseline is stable, layer complexity gradually and verify behavior after each change. This staged workflow is more predictable than changing multiple variables at once and trying to infer root cause afterward.

2) Use tf.data with a generator and parallel map transforms

python
1import tensorflow as tf
2import numpy as np
3
4def sample_generator(num_samples=100000):
5    for _ in range(num_samples):
6        x = np.random.uniform(-1.0, 1.0, size=(64, 64, 3)).astype(np.float32)
7        y = np.float32(x.mean() > 0.0)
8        yield x, y
9
10output_signature = (
11    tf.TensorSpec(shape=(64, 64, 3), dtype=tf.float32),
12    tf.TensorSpec(shape=(), dtype=tf.float32),
13)
14
15ds = tf.data.Dataset.from_generator(sample_generator, output_signature=output_signature)
16ds = ds.map(lambda x, y: (tf.image.random_flip_left_right(x), y),
17            num_parallel_calls=tf.data.AUTOTUNE)
18ds = ds.shuffle(4096).batch(64).prefetch(tf.data.AUTOTUNE)

This baseline snippet is intentionally conservative. It prioritizes readability, deterministic behavior, and explicit control points over clever shortcuts. For production, you can tune performance later, but first ensure the pipeline is correct and repeatable. If this step does not behave as expected, freeze further refactors and diagnose here; debugging gets exponentially harder once additional abstractions are layered on top.

3) Add caching, profiling, and deterministic modes for debugging

python
1options = tf.data.Options()
2options.experimental_deterministic = False  # True for reproducible debugging
3
4ds = ds.with_options(options)
5
6def normalize(x, y):
7    x = tf.clip_by_value(x, -1.0, 1.0)
8    return x, y
9
10ds = ds.map(normalize, num_parallel_calls=tf.data.AUTOTUNE)
11
12for step, (xb, yb) in enumerate(ds.take(3)):
13    tf.print("batch", step, "shape", tf.shape(xb), "label mean", tf.reduce_mean(yb))

Operational guardrails are what turn a working demo into a maintainable system. Add logging around key transitions, monitor latency and error classes, and define clear retry or fallback policy where failures are expected. Avoid silent recovery paths that hide data quality or state issues. Instead, emit structured signals that make post-incident analysis straightforward.

4) Validate behavior with repeatable checks

Run a short training loop with fixed seeds and inspect class balance, tensor dtypes, and batch latency. A good pipeline keeps batch production consistently faster than model step time. Write a short verification checklist that can run in local development, CI, and pre-release environments. Include both success-path assertions and at least one intentional failure case. Over time, this checklist becomes regression protection: it documents assumptions, catches environment drift, and prevents future edits from reintroducing the same class of bug.

Common Pitfalls

  • Doing heavy augmentation in pure Python loops instead of TensorFlow ops, which throttles throughput.
  • Applying random transforms that change labels without updating the target generation logic.
  • Skipping prefetch, causing the accelerator to wait for input batches.
  • Using unbounded shuffle buffers that exceed memory on long-running jobs.
  • Assuming reproducibility while deterministic mode is disabled and seeds are unset.

Summary

On-the-fly data generation works best when the pipeline is vectorized, measurable, and explicit about randomness and label integrity. The key pattern is consistent across stacks: keep the core path simple, instrument the edges, and validate with deterministic tests before scaling complexity.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track 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.

Practice ML system design

All Rights Reserved.