TopolearnA slow input pipeline leaves expensive GPUs idle; tuning tf.data is often the cheapest speedup available.
Training steps on a GPU are fast; if loading and preprocessing the next batch takes longer, the accelerator waits. tf.data has built-in tools to overlap and parallelise that work:
.map(fn, num_parallel_calls=tf.data.AUTOTUNE)runs preprocessing on several CPU cores..prefetch(tf.data.AUTOTUNE)at the end of the pipeline prepares the next batches while the model trains on the current one. Almost every pipeline should end with it..cache()stores elements after the first epoch, in memory or in a file (cache("path")). Put it after expensive deterministic work but before random augmentation and shuffling..interleave(...)reads from several files concurrently, essential for sharded TFRecord datasets on network storage.- Vectorise: batching first and then mapping a function that handles a whole batch is often faster than mapping per element.
AUTOTUNE lets the runtime pick buffer sizes and parallelism dynamically.
A typical order: list files, interleave reads, parse/decode with parallel map, cache, shuffle, batch, augment, prefetch. To diagnose, time an epoch of the dataset alone without the model, or use the TensorBoard profiler's input-pipeline analyser. If the dataset alone is slower than training, the pipeline is your bottleneck.
Example
import tensorflow as tf
AUTOTUNE = tf.data.AUTOTUNE
def load(path):
img = tf.io.decode_jpeg(tf.io.read_file(path), channels=3)
return tf.image.resize(img, [224, 224]) / 255.0
def augment(batch):
return tf.image.random_flip_left_right(batch)
files = tf.data.Dataset.list_files("images/*.jpg", shuffle=True)
ds = (files
.map(load, num_parallel_calls=AUTOTUNE) # decode on several cores
.cache("/tmp/img_cache") # reuse decoded images after epoch 1
.shuffle(1000)
.batch(64)
.map(augment, num_parallel_calls=AUTOTUNE) # random work after the cache
.prefetch(AUTOTUNE)) # overlap input with training
# model.fit(ds, epochs=10)Remember
- Parallel
map+ finalprefetch(AUTOTUNE)in almost every pipeline - Cache after expensive deterministic steps, before random ones
- Time the dataset alone to find input bottlenecks
Watch out for
- Caching after random augmentation, which freezes one set of augmentations for all epochs.