tf.data
tf.data.Dataset is a high-throughput input pipeline. Lazy, parallel, prefetched — feeds the GPU at line speed. The unsexy but critical layer of any production training loop.
Build, transform, prefetch, files
EXAMPLE
import tensorflow as tf
# 1) From a Python list / array
ds = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5])
for x in ds:
print(x.numpy())
# Tuple of (features, labels)
X = tf.random.normal((1000, 20))
y = tf.random.uniform((1000,), maxval=10, dtype=tf.int32)
ds = tf.data.Dataset.from_tensor_slices((X, y))
# 2) Transform with .map
ds = ds.map(lambda x, y: (x / 255.0, y), num_parallel_calls=tf.data.AUTOTUNE)
# 3) Shuffle, batch, prefetch — the standard finishing combo
ds = (ds
.shuffle(buffer_size=10_000)
.batch(64)
.prefetch(tf.data.AUTOTUNE))
# 4) Read from files — image folders
train_ds = tf.keras.utils.image_dataset_from_directory(
'data/train',
image_size = (224, 224),
batch_size = 32,
label_mode = 'categorical',
shuffle = True,
seed = 42,
)
val_ds = tf.keras.utils.image_dataset_from_directory(
'data/val',
image_size = (224, 224),
batch_size = 32,
)
class_names = train_ds.class_names
# 5) Read text files
ds = tf.data.TextLineDataset(['log1.txt', 'log2.txt'])
# CSV with auto-parsing
ds = tf.data.experimental.make_csv_dataset(
'users.csv',
batch_size = 32,
label_name = 'class',
num_epochs = 1,
shuffle = True,
)
# 6) TFRecord — the high-throughput format
def parse(example):
features = {
'image': tf.io.FixedLenFeature([], tf.string),
'label': tf.io.FixedLenFeature([], tf.int64),
}
p = tf.io.parse_single_example(example, features)
img = tf.image.decode_jpeg(p['image'], channels=3)
img = tf.image.resize(img, [224, 224]) / 255.0
return img, p['label']
ds = (tf.data.TFRecordDataset(['train-001.tfrecord', 'train-002.tfrecord'])
.map(parse, num_parallel_calls=tf.data.AUTOTUNE)
.shuffle(10_000)
.batch(64)
.prefetch(tf.data.AUTOTUNE))
# 7) Pipeline composition — the order matters
ds = (raw
.cache() # cache after read, before augment
.shuffle(10_000)
.map(parse, num_parallel_calls=tf.data.AUTOTUNE)
.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
.batch(64)
.prefetch(tf.data.AUTOTUNE))
# Rules of thumb:
# - cache: after expensive deterministic work (file decoding) and BEFORE random augment
# - shuffle: BEFORE batch (shuffles examples, not batches)
# - prefetch: ALWAYS at the end
# - batch: usually right before prefetch
# - num_parallel_calls=AUTOTUNE on map: let TF decide threads
# 8) Augmentation — vectorise across the batch
def augment(x, y):
x = tf.image.random_flip_left_right(x)
x = tf.image.random_brightness(x, 0.1)
x = tf.image.random_contrast(x, 0.9, 1.1)
return x, y
# Or use Keras preprocessing layers in-graph (run on GPU):
aug = tf.keras.Sequential([
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.1),
tf.keras.layers.RandomZoom(0.1),
])
ds = ds.map(lambda x, y: (aug(x, training=True), y), num_parallel_calls=tf.data.AUTOTUNE)
# 9) Generators — when data is dynamic
def gen():
for i in range(1000):
yield (np.random.randn(20).astype('float32'), np.random.randint(10))
ds = tf.data.Dataset.from_generator(
gen,
output_signature=(
tf.TensorSpec(shape=(20,), dtype=tf.float32),
tf.TensorSpec(shape=(), dtype=tf.int32),
),
)
# 10) Zip / interleave — combine multiple sources
ds1 = tf.data.Dataset.range(10)
ds2 = tf.data.Dataset.range(100, 110)
zipped = tf.data.Dataset.zip((ds1, ds2)) # (0, 100), (1, 101), ...
images = tf.data.Dataset.list_files('data/*.jpg')
features = (images
.interleave(
lambda fn: tf.data.TFRecordDataset(fn),
cycle_length=4,
num_parallel_calls=tf.data.AUTOTUNE,
))
# 11) Caching — disk vs memory
ds = ds.cache() # in-memory (fits)
ds = ds.cache('/tmp/cache') # disk-backed (huge dataset)
# 12) Sharding — distributed training
ds = ds.shard(num_shards=4, index=0)
# Each worker reads its 1/4 slice
# 13) Inspect the pipeline
for batch_x, batch_y in ds.take(1):
print(batch_x.shape, batch_y.shape)
# Display options to debug
print(ds.element_spec)
print(ds.cardinality()) # known size if applicable
# 14) Performance — measure not guess
import time
start = time.time()
for _ in ds.take(100): pass
print(f'{(time.time()-start):.2f}s for 100 batches')
# TF tools:
# tf.profiler — profile bottlenecks (CPU bound? IO bound? GPU starvation?)
# tf.data.experimental.AUTOTUNE — let TF tune threads / buffer sizes
# 15) Common bugs
# • Shuffle AFTER batch → only shuffles batch order; useless
# • Augment BEFORE cache → caches augmented data; loses variety
# • No prefetch → GPU idles waiting for next batch
# • Loading whole dataset to memory → OOM (use .cache() carefully)
# • num_parallel_calls=1 → underuses CPU; use AUTOTUNE
# 16) Best practices
# • Always end with prefetch(AUTOTUNE)
# • Use TFRecord for big training datasets — fastest input format
# • Read .repeat() before .batch() for multi-epoch training (or just pass epochs to fit)
# • Profile with TensorBoard's Profile tab → identify GPU starvation
Why it matters
A great tf.data pipeline keeps your GPU at 100% utilisation. Bad pipelines leave the GPU idle 80% of the time. Order: parse → cache → shuffle → augment → batch → prefetch.
Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.
Example
Example
import tensorflow as tf
ds = (tf.data.Dataset
.from_tensor_slices((X, y))
.shuffle(1000)
.batch(64)
.prefetch(tf.data.AUTOTUNE))
Try it Yourself »
Exercise
Build a streaming dataset.
ds = tf.
.Dataset.from_tensor_slices(arr)
Four letters.
Discussion
Loading…