Data Pipeline

Code

import tensorflow as tf

# From tensors
ds = tf.data.Dataset.from_tensor_slices(
    (tf.random.normal((100, 32)), tf.random.uniform((100,), 0, 2, tf.int32))
)

# Transform pipeline
ds = (ds
      .shuffle(buffer_size=1000)
      .batch(32)
      .map(lambda x, y: (x, tf.one_hot(y, depth=2)),
           num_parallel_calls=tf.data.AUTOTUNE)
      .prefetch(tf.data.AUTOTUNE))

for x, y in ds.take(1):
    print(x.shape, y.shape)

# Read images from disk
def load_image(path, label):
    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, [224, 224])
    return img, label

paths = tf.data.Dataset.from_tensor_slices(["a.jpg", "b.jpg"])
labels = tf.data.Dataset.from_tensor_slices([0, 1])
image_ds = tf.data.Dataset.zip((paths, labels)).map(load_image).batch(8)

Explanation

tf.data.Dataset chains lazy transformations like map, shuffle, batch, and prefetch to build a streaming pipeline. prefetch(AUTOTUNE) overlaps data preparation with training, hiding I/O latency. map with num_parallel_calls parallelizes preprocessing, which is critical for image pipelines reading from disk.

More TensorFlow Snippets