Prefetching for performance in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Prefetching for performance

Problem:You have a TensorFlow model training on image data using tf.data.Dataset. The training is slow because the data loading and preprocessing block the GPU from running efficiently.

Current Metrics:Training time per epoch: 120 seconds; Training accuracy after 5 epochs: 75%; Validation accuracy after 5 epochs: 72%

Issue:The model training is slow due to input pipeline bottleneck. The GPU waits for data because the dataset does not use prefetching.

Your Task

Add prefetching to the TensorFlow data pipeline to reduce input latency and improve training speed without changing the model architecture or batch size.

Do not change the model architecture.

Keep batch size and number of epochs the same.

Only modify the data pipeline to add prefetching.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf

# Load example dataset
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# Normalize images
train_images = train_images / 255.0

# Create tf.data.Dataset
train_ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_ds = train_ds.shuffle(10000).batch(64)

# Add prefetching to improve performance
train_ds = train_ds.prefetch(tf.data.AUTOTUNE)

# Define a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])

# Compile model
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# Train model
history = model.fit(train_ds, epochs=5, validation_data=(test_images / 255.0, test_labels))

Added .prefetch(tf.data.AUTOTUNE) to the tf.data.Dataset pipeline to allow data loading and preprocessing to happen in parallel with model training.

Results Interpretation

Before adding prefetching: Training time per epoch was 120 seconds with 75% training accuracy and 72% validation accuracy.

After adding prefetching: Training time per epoch reduced to 80 seconds while maintaining the same accuracy levels.

Prefetching allows the data pipeline to prepare the next batch of data while the model is training on the current batch. This reduces idle GPU time and speeds up training without affecting model accuracy.

Bonus Experiment

Try adding data caching to the pipeline along with prefetching to see if training speed improves further.

💡 Hint

Use the .cache() method before prefetching to keep data in memory if it fits.