tf.data input pipelines

TensorFlow · Essentials · 12 min

tf.data is how TensorFlow streams data into training efficiently, from in-memory arrays to terabytes of files.

A `tf.data.Dataset` is a sequence of elements (tensors or tuples and dicts of tensors) with chainable transformations. Each method returns a new dataset:

  • Dataset.from_tensor_slices((x, y)) slices arrays along the first axis into (x_i, y_i) pairs.
  • .map(fn) transforms each element; fn is traced into a graph, so use TF ops inside it.
  • .shuffle(buffer_size) shuffles using a buffer; a buffer at least as large as the dataset gives a full shuffle.
  • .batch(n) groups elements; the last batch may be smaller unless drop_remainder=True.
  • .take(n), .skip(n), .filter(pred), .repeat() do what their names say.

Datasets are Python iterables in eager mode, so for xb, yb in ds: works, and model.fit(ds) accepts them directly (do not also pass batch_size).

ds.element_spec shows the structure, shapes and dtypes of each element, handy for debugging. Other sources include tf.data.TextLineDataset, TFRecordDataset and Dataset.from_generator for arbitrary Python generators.

The order of transformations matters: shuffle before batching to shuffle examples (not whole batches), and map before batching when the function processes single examples.

Example

import tensorflow as tf

x = tf.range(10, dtype=tf.float32)
ds = tf.data.Dataset.from_tensor_slices((x, x ** 2))
ds = (ds.map(lambda a, b: (a / 10, b))      # per-example transform
        .shuffle(buffer_size=10)            # full shuffle for 10 items
        .batch(4))

for xb, yb in ds:
    print(xb.shape, yb.shape)
# (4,) (4,)
# (4,) (4,)
# (2,) (2,)

print(ds.element_spec[0])   # TensorSpec(shape=(None,), dtype=tf.float32, name=None)
print(len(list(ds.unbatch())))   # 10

Remember

  • from_tensor_slices + map + shuffle + batch
  • Shuffle before batch; map single examples before batch
  • Pass datasets straight to model.fit

Watch out for

  • Using a tiny shuffle buffer on sorted data, which barely shuffles at all.

Check yourself

  1. ds = Dataset.from_tensor_slices(range(10)).batch(3). How many batches does it yield?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: www.tensorflow.org/guide/data