tf.data.Dataset is Keras’s answer to Dataset + DataLoader — batching, shuffling, and prefetching, but as a chained, lazy pipeline. Feed it straight into fit().
Why: where PyTorch splits data handling into a Dataset (one item) and a DataLoader (batching/shuffling), TensorFlow combines both into tf.data.Dataset — a lazy, chainable pipeline of transformations. When: build a Dataset, then chain shuffle/batch/prefetch, and pass it to fit(). Where: prefetch is the tf.data equivalent of DataLoader’s num_workers/pin_memory — it overlaps data prep with training so the GPU never waits.
import tensorflow as tf
ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)) # like TensorDataset
train_ds = (ds
.shuffle(1000) # DataLoader(shuffle=True)
.batch(32) # DataLoader(batch_size=32)
.prefetch(tf.data.AUTOTUNE)) # like num_workers/pin_memory — keep the GPU fed
model.fit(train_ds, epochs=10) # pass the Dataset straight inWhy: preprocessing and augmentation attach to the pipeline with .map(), which runs lazily and in parallel — the tf.data analog of transforms in a PyTorch Dataset’s __getitem__. When: map normalization or augmentation onto the training pipeline (apply augmentation to train only, exactly as in PyTorch). Where: num_parallel_calls=AUTOTUNE parallelizes the mapping, similar to multiple DataLoader workers.
import tensorflow as tf
def preprocess(x, y):
x = tf.cast(x, tf.float32) / 255.0 # normalize, like a transform
return x, y
train_ds = (tf.data.Dataset.from_tensor_slices((images, labels))
.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE) # parallel, lazy
.shuffle(1000).batch(64).prefetch(tf.data.AUTOTUNE))Why: for images on disk, Keras offers helpers that build a batched, labeled tf.data pipeline in one call — the equivalent of writing a custom PyTorch Dataset plus torchvision’s ImageFolder. When: use image_dataset_from_directory for folder-structured image data; use tf.data readers for other sources. Where: the result is a Dataset you pass straight to fit(), so the rest of your training code is unchanged.
from tensorflow import keras
# Folder-per-class images -> a batched, labeled pipeline (like ImageFolder):
train_ds = keras.utils.image_dataset_from_directory(
"data/train", image_size=(180, 180), batch_size=32, label_mode="int",
)
val_ds = keras.utils.image_dataset_from_directory(
"data/val", image_size=(180, 180), batch_size=32, label_mode="int",
)
model.fit(train_ds, validation_data=val_ds, epochs=10)