How to optimize TensorFlow data pipeline performance

A slow tf.data pipeline can leave the accelerator or training loop waiting for the next batch even when the model code is ready to run. The usual fix is to move deterministic input work into a pipeline that can process examples in parallel, reuse prepared data, and queue future batches before the current batch finishes training.

TensorFlow exposes those controls through Dataset.map(), cache(), shuffle(), batch(), and prefetch(). tf.data.AUTOTUNE lets the runtime choose parallelism and prefetch buffer sizes, while cache() stores the output of earlier transformations after the first full pass through the dataset.

The cache boundary decides which work is reused and which work stays fresh. Put deterministic parsing, decoding, resizing, or normalization before cache() when the prepared examples fit in memory or local storage, keep training-only randomization after that boundary, and benchmark the real pipeline because CPU, storage, and transformation cost determine the actual speedup.

1. Open the Python environment that runs the training job.

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

   Use the same environment that runs model.fit() so the benchmark uses the same TensorFlow package and CPU thread behavior.
   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Save the benchmark as

   optimize-tfdata-pipeline.py

   .

   optimize-tfdata-pipeline.py

   import os
   import time
    
   import numpy as np
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
    
   tf.get_logger().setLevel("ERROR")
   tf.keras.utils.set_random_seed(2026)
    
   AUTOTUNE = tf.data.AUTOTUNE
   EXAMPLES = 640
   FEATURES = 3
   BATCH_SIZE = 32
   EPOCHS = 2
   DELAY_SECONDS = 0.002
   FEATURE_SCALE = np.array([1000.0, 100.0, 10.0], dtype=np.float32)
    
    
   features = tf.random.uniform(
       (EXAMPLES, FEATURES),
       minval=0.0,
       maxval=100.0,
       dtype=tf.float32,
       seed=7,
   )
   labels = tf.cast(tf.reduce_sum(features, axis=1) > 150.0, tf.float32)
    
    
   def pause_and_normalize(value):
       time.sleep(DELAY_SECONDS)
       return (value.astype(np.float32) / FEATURE_SCALE).astype(np.float32)
    
    
   def slow_preprocess(feature_row, label):
       normalized = tf.numpy_function(pause_and_normalize, [feature_row], tf.float32)
       normalized.set_shape([FEATURES])
       return normalized, label
    
    
   def build_pipeline(optimized):
       dataset = tf.data.Dataset.from_tensor_slices((features, labels))
       if optimized:
           dataset = dataset.map(slow_preprocess, num_parallel_calls=AUTOTUNE)
           dataset = dataset.cache()
           dataset = dataset.shuffle(
               buffer_size=EXAMPLES,
               seed=2026,
               reshuffle_each_iteration=False,
           )
           dataset = dataset.batch(BATCH_SIZE)
           dataset = dataset.prefetch(AUTOTUNE)
           return dataset
    
       dataset = dataset.map(slow_preprocess)
       dataset = dataset.shuffle(
           buffer_size=EXAMPLES,
           seed=2026,
           reshuffle_each_iteration=False,
       )
       dataset = dataset.batch(BATCH_SIZE)
       return dataset
    
    
   def consume(dataset):
       batches = 0
       checksum = 0.0
       started = time.perf_counter()
       for _ in range(EPOCHS):
           for batch_features, _ in dataset:
               checksum += float(tf.reduce_sum(batch_features).numpy())
               batches += 1
       return time.perf_counter() - started, batches, checksum
    
    
   baseline_seconds, baseline_batches, baseline_checksum = consume(build_pipeline(False))
   optimized_dataset = build_pipeline(True)
   optimized_seconds, optimized_batches, optimized_checksum = consume(optimized_dataset)
   batch_features, batch_labels = next(iter(optimized_dataset))
   speedup = baseline_seconds / optimized_seconds
    
   print(f"TensorFlow {tf.__version__}")
   print(f"examples={EXAMPLES}")
   print(f"batch_size={BATCH_SIZE}")
   print(f"baseline_seconds={baseline_seconds:.3f}")
   print(f"optimized_seconds={optimized_seconds:.3f}")
   print(f"speedup={speedup:.2f}x")
   print(f"optimized_batches={optimized_batches}")
   print(f"train_batch_shape={tuple(batch_features.shape)}")
   print(f"train_label_shape={tuple(batch_labels.shape)}")
   print(f"checksums_match={abs(baseline_checksum - optimized_checksum) < 1e-3}")
   print(f"optimized_dataset={type(optimized_dataset).__name__}")

   The synthetic delay stands in for deterministic preprocessing work such as parsing, decoding, resizing, or numeric normalization. Replace slow_preprocess() with the project preprocessing function when moving the pattern into real training code.

3. Run the benchmark and confirm the optimized pipeline is faster while the checksums and batch shapes still match.

   $ python3 optimize-tfdata-pipeline.py
   TensorFlow 2.21.0
   examples=640
   batch_size=32
   baseline_seconds=4.814
   optimized_seconds=0.269
   speedup=17.87x
   optimized_batches=40
   train_batch_shape=(32, 3)
   train_label_shape=(32,)
   checksums_match=True
   optimized_dataset=_PrefetchDataset

   checksums_match=True confirms that the optimized pipeline produced the same normalized examples as the baseline path. The exact timings vary by machine, but optimized_seconds should fall below baseline_seconds when the input stage has parallelizable work.

4. Apply the optimized order to the training dataset.

   AUTOTUNE = tf.data.AUTOTUNE
    
   train_ds = (
       raw_train_ds
       .map(parse_and_preprocess, num_parallel_calls=AUTOTUNE)
       .cache()
       .shuffle(buffer_size=10000, reshuffle_each_iteration=True)
       .batch(64)
       .prefetch(AUTOTUNE)
   )

   Place cache() after deterministic preprocessing and before the training shuffle so later epochs reuse prepared examples while each epoch can still receive a fresh order.

   Do not cache after random augmentation, random cropping, or other training-only randomness unless the goal is to freeze one generated version of each example. Random image pipelines should cache the deterministic base data before augmentation.
   Related: How to build an image augmentation pipeline in TensorFlow

5. Keep validation and test datasets deterministic.

   validation_ds = (
       raw_validation_ds
       .map(parse_and_preprocess, num_parallel_calls=AUTOTUNE)
       .cache()
       .batch(64)
       .prefetch(AUTOTUNE)
   )

   Validation and test input should normally skip shuffle() and random augmentation so metric changes come from the model, not from a moving evaluation input.

6. Use local storage cache only when the prepared data no longer fits comfortably in memory.

   train_ds = (
       raw_train_ds
       .map(parse_and_preprocess, num_parallel_calls=AUTOTUNE)
       .cache("/tmp/tfdata-train-cache")
       .shuffle(buffer_size=10000)
       .batch(64)
       .prefetch(AUTOTUNE)
   )

   A stale cache file can hide preprocessing changes. Delete or rename the cache path after changing parsing, resizing, feature order, label handling, or augmentation boundaries.

7. Measure the real training pipeline with a fixed batch count before changing model code.

   import time
    
   started = time.perf_counter()
   for _ in train_ds.take(100):
       pass
    
   print(f"100_batches_seconds={time.perf_counter() - started:.3f}")

   Run the same check before and after the pipeline change, then profile model.fit() if the input stage still dominates training time.
   Related: How to profile TensorFlow training in TensorBoard

8. Remove the temporary benchmark after the project pipeline reports the expected shapes and faster batch delivery.

   $ rm optimize-tfdata-pipeline.py