Without prefetching, a training loop processes data strictly sequentially: prepare one batch, then train on it, then prepare the next batch, then train on that, with the model sitting idle during every data-preparation step. prefetch() overlaps these two stages instead, preparing the next batch, or batches, on the CPU while the model is still busy training on the current one, so the two stages run concurrently rather than one after another. Passing buffer_size=tf.data.AUTOTUNE lets TensorFlow automatically tune how many batches to prefetch based on runtime conditions, and prefetch() is typically the very last transformation applied in a tf.data pipeline.
1Understanding dataset.prefetch()
Without prefetching, a training loop processes data strictly sequentially: prepare one batch, then train on it, then prepare the next batch, then train on that, with the model sitting idle during every data-preparation step. prefetch() overlaps these two stages instead, preparing the next batch, or batches, on the CPU while the model is still busy training on the current one, so the two stages run concurrently rather than one after another. Passing buffer_size=tf.data.AUTOTUNE lets TensorFlow automatically tune how many batches to prefetch based on runtime conditions, and prefetch() is typically the very last transformation applied in a tf.data pipeline.
Place .prefetch(tf.data.AUTOTUNE) as the last step of a tf.data pipeline, after batch() — since it overlaps data preparation with training, it usually provides a meaningful speedup for essentially free, at the cost of a small, tunable amount of extra memory.
import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])
pipeline = dataset.batch(2).prefetch(tf.data.AUTOTUNE)
for element in pipeline:
print(element.numpy())2Practical Example
Here is a real-world application of dataset.prefetch() showing how it is used in production TensorFlow code.
import tensorflow as tf
dataset = tf.data.Dataset.range(100)
pipeline = dataset.batch(10).prefetch(buffer_size=2)
print(sum(1 for _ in pipeline))3Best Practices
Follow these guidelines when working with dataset.prefetch():
1. Add .prefetch(tf.data.AUTOTUNE) as the final step of every tf.data pipeline used for training, to overlap data preparation with model computation
2. Let AUTOTUNE choose the prefetch buffer size automatically rather than guessing a fixed number, since the ideal amount depends on runtime hardware and pipeline characteristics
3. Combine prefetch() with num_parallel_calls=AUTOTUNE on map() for a pipeline that both parallelizes preprocessing and overlaps it with training
Tip: Place .prefetch(tf.data.AUTOTUNE) as the last step of a tf.data pipeline, after batch() — since it overlaps data preparation with training, it usually provides a meaningful speedup for essentially free, at the cost of a small, tunable amount of extra memory.
import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])
pipeline = dataset.batch(2).prefetch(tf.data.AUTOTUNE)
for element in pipeline:
print(element.numpy())