Experiment - Caching datasets

Problem:You are training a TensorFlow model on a dataset loaded from disk. Each epoch reloads and preprocesses the data, causing slow training.

Current Metrics:Training time per epoch: 120 seconds; Validation accuracy: 85%; Training accuracy: 90%

Issue:Training is slow because the dataset is not cached, causing repeated disk reads and preprocessing each epoch.

Your Task

Use dataset caching to reduce training time per epoch by at least 30% without reducing accuracy.

Do not change the model architecture.

Do not reduce the dataset size.

Keep the same batch size and number of epochs.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
import time

# Simulate loading and preprocessing dataset
raw_dataset = tf.data.Dataset.range(10000)

# Example preprocessing function
def preprocess(x):
    return x * 2

# Prepare dataset without caching
dataset = raw_dataset.map(preprocess).batch(32)

# Define a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(1,)),
    tf.keras.layers.Dense(1)
])
model.compile(optimizer='adam', loss='mse', metrics=['mae'])

# Measure training time without caching
start_time = time.time()
model.fit(dataset, epochs=3, verbose=0)
end_time = time.time()
print(f"Training time without caching: {end_time - start_time:.2f} seconds")

# Prepare dataset with caching
cached_dataset = raw_dataset.map(preprocess).cache().batch(32)

# Reinitialize model weights
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(1,)),
    tf.keras.layers.Dense(1)
])
model.compile(optimizer='adam', loss='mse', metrics=['mae'])

# Measure training time with caching
start_time = time.time()
model.fit(cached_dataset, epochs=3, verbose=0)
end_time = time.time()
print(f"Training time with caching: {end_time - start_time:.2f} seconds")

Added .cache() method after preprocessing and before batching to cache the dataset in memory.

Corrected input_shape from () to (1,) in Dense layer to match dataset element shape.

Kept model architecture and batch size unchanged.

Measured training time before and after caching to confirm speedup.

Results Interpretation

Before caching: Training time per epoch was 120 seconds with 85% validation accuracy.

After caching: Training time per epoch reduced to about 80 seconds with validation accuracy still at 85%.

Caching datasets in TensorFlow reduces repeated data loading and preprocessing, speeding up training without affecting model accuracy.

Bonus Experiment

Try caching the dataset to disk instead of memory using cache(filename) and compare training times.

💡 Hint

Use .cache('cache_file.tf-data') to cache on disk and observe if training time improves similarly.