How to build an image augmentation pipeline in TensorFlow

Building an image augmentation pipeline in TensorFlow lets the training input stage generate fresh image variants without rewriting the source files on disk. This keeps the model exposed to flips, rotations, and zoom changes while the labels and dataset split stay stable.

A tf.data.Dataset pipeline separates deterministic work such as resizing and rescaling from training-only random transforms. Current Keras preprocessing layers such as RandomFlip, RandomRotation, and RandomZoom can run inside a mapped dataset batch when they are called with training=True, and prefetch() keeps those augmented batches queued for the training loop.

The cache boundary matters in this workflow. If cache() is placed after the random augmentation stage, later epochs reuse the same transformed pixels instead of generating new variants, while validation and test datasets should stay on the deterministic preprocessing path so their metrics measure the real model state instead of a moving random target.

1. Open a terminal in a Python environment where TensorFlow already imports cleanly.

   $ python3 -c "import tensorflow as tf; print(tf.__version__); print(tf.keras.__version__)"
   2.21.0
   3.15.0

   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Save the demo as

   image-augmentation-pipeline.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   layers = tf.keras.layers
    
   tf.get_logger().setLevel("ERROR")
   tf.keras.utils.set_random_seed(7)
    
   IMAGE_SIZE = 64
   BATCH_SIZE = 4
    
   def make_image(index):
       index = tf.cast(index, tf.float32)
       grid = tf.linspace(0.0, 1.0, 48)
       xx, yy = tf.meshgrid(grid, grid)
       xx = tf.expand_dims(xx, -1)
       yy = tf.expand_dims(yy, -1)
       diagonal = tf.cast(xx > yy, tf.float32)
       image = tf.concat(
           [
               tf.clip_by_value(xx + index * 0.04, 0.0, 1.0),
               tf.clip_by_value(1.0 - yy * 0.8, 0.0, 1.0),
               tf.clip_by_value(diagonal * 0.7 + index * 0.02, 0.0, 1.0),
           ],
           axis=-1,
       )
       return tf.cast(image * 255.0, tf.float32)
    
   images = tf.stack([make_image(i) for i in range(8)], axis=0)
   labels = tf.constant([0, 1, 0, 1, 0, 1, 0, 1], dtype=tf.int32)
    
   resize_and_rescale = tf.keras.Sequential(
       [
           layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),
           layers.Rescaling(1.0 / 255.0),
       ]
   )
    
   augment = tf.keras.Sequential(
       [
           layers.RandomFlip("horizontal", seed=7),
           layers.RandomRotation(0.15, seed=7),
           layers.RandomZoom(0.1, seed=7),
       ]
   )
    
   def prepare(image, label):
       image = resize_and_rescale(image)
       return image, label
    
   base_ds = (
       tf.data.Dataset.from_tensor_slices((images, labels))
       .map(prepare, num_parallel_calls=tf.data.AUTOTUNE)
       .cache()
   )
    
   train_ds = (
       base_ds
       .shuffle(len(labels), seed=7, reshuffle_each_iteration=False)
       .batch(BATCH_SIZE)
   )
    
   train_aug_ds = train_ds.map(
       lambda x, y: (augment(x, training=True), y),
       num_parallel_calls=tf.data.AUTOTUNE,
   ).prefetch(tf.data.AUTOTUNE)
    
   val_ds = base_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
    
   first_train_images, first_train_labels = next(iter(train_aug_ds))
   second_train_images, _ = next(iter(train_aug_ds))
   val_images, _ = next(iter(val_ds))
   base_images, _ = next(iter(base_ds.batch(BATCH_SIZE)))
    
   pixel_delta = float(tf.reduce_sum(tf.abs(first_train_images - second_train_images)).numpy())
   val_matches_base = bool(tf.reduce_all(tf.abs(val_images - base_images) < 1e-6).numpy())
    
   print(f"TensorFlow {tf.__version__}")
   print(f"Keras {tf.keras.__version__}")
   print(f"train_batch_shape={tuple(first_train_images.shape)}")
   print(f"train_labels={first_train_labels.numpy().tolist()}")
   print(f"augmentation_changes_pixels={pixel_delta > 0.0}")
   print(f"validation_matches_base={val_matches_base}")
   print(
       "pixel_value_range="
       f"({float(tf.reduce_min(first_train_images).numpy()):.3f}, "
       f"{float(tf.reduce_max(first_train_images).numpy()):.3f})"
   )

   The demo fixes reshuffle_each_iteration=False so the comparison proves the augmentation effect instead of a different batch order. Restore the default reshuffle behavior in real training when each epoch should also see a new image order.
   Related: How to set a random seed in TensorFlow

3. Run the script and confirm that the training batch changes between augmentation passes while the validation batch still matches the cached base pipeline.

   $ python3 image-augmentation-pipeline.py
   TensorFlow 2.21.0
   Keras 3.15.0
   train_batch_shape=(4, 64, 64, 3)
   train_labels=[0, 1, 1, 0]
   augmentation_changes_pixels=True
   validation_matches_base=True
   pixel_value_range=(0.000, 1.000)

   augmentation_changes_pixels=True shows that the training path applies random transforms, and validation_matches_base=True shows that the validation path keeps the cached deterministic images.

   Related: How to optimize TensorFlow data pipeline performance

4. Reuse the same cache-then-augment structure in the real training pipeline so only the training batches receive random transforms.

   AUTOTUNE = tf.data.AUTOTUNE
    
   train_ds = (
       train_raw_ds
       .map(lambda x, y: (resize_and_rescale(x), y), num_parallel_calls=AUTOTUNE)
       .cache()
       .map(lambda x, y: (augment(x, training=True), y), num_parallel_calls=AUTOTUNE)
       .prefetch(AUTOTUNE)
   )
    
   val_ds = (
       val_raw_ds
       .map(lambda x, y: (resize_and_rescale(x), y), num_parallel_calls=AUTOTUNE)
       .cache()
       .prefetch(AUTOTUNE)
   )
    
   model.fit(train_ds, validation_data=val_ds, epochs=10)

   Do not cache after the random augmentation map unless the goal is to freeze one augmented copy of each image. Caching after augmentation stops later epochs from seeing new flips, rotations, or zoom values.

   Related: How to train, evaluate, and run prediction with a Keras model
   Related: How to run transfer learning in Keras

5. Remove the temporary demo script after the augmentation pipeline check passes.

   $ rm image-augmentation-pipeline.py

• Call the augmentation layer with training=True inside Dataset.map() or the random image layers stay in inference mode and return the unmodified batch.
• Keep validation and test pipelines deterministic so evaluation compares models, not different random image variants from run to run.
• The sample uses synthetic images only to prove the pipeline shape; replace the source dataset with the actual training batches after the control flow is working.