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;fnis 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 unlessdrop_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()))) # 10Remember
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.