iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

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)

Discussion

Loading…