parallelising tf.data.Dataset.from_generator

TensorFlow

tf.data

parallel processing

data pipeline

Dataset generation

parallelising tf.data.Dataset.from_generator

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

TensorFlow & Keras Concurrency

Introduction

tf.data.Dataset.from_generator runs a Python generator in a single thread, making it a bottleneck in training pipelines. To parallelize it, use tf.data.Dataset.interleave with multiple generator instances, or restructure your pipeline to use tf.data.Dataset.map with num_parallel_calls for the heavy processing after the generator yields file paths or indices. For maximum performance, convert your data to TFRecord format and use tf.data.TFRecordDataset with parallel reads instead of generators.

The Problem

python

1import tensorflow as tf
2import numpy as np
3
4def data_generator():
5    for i in range(10000):
6        # Slow preprocessing (I/O, augmentation, etc.)
7        image = load_and_preprocess(f"image_{i}.jpg")
8        label = get_label(i)
9        yield image, label
10
11dataset = tf.data.Dataset.from_generator(
12    data_generator,
13    output_signature=(
14        tf.TensorSpec(shape=(224, 224, 3), dtype=tf.float32),
15        tf.TensorSpec(shape=(), dtype=tf.int32),
16    )
17)

The generator runs in the Python GIL, producing one sample at a time. Even with .prefetch(), the generator itself cannot be parallelized directly.

Method 1: Interleave Multiple Generators

Create multiple generator instances and interleave their outputs:

python

1def make_generator(shard_id, num_shards):
2    def generator():
3        for i in range(shard_id, 10000, num_shards):
4            image = load_and_preprocess(f"image_{i}.jpg")
5            label = get_label(i)
6            yield image, label
7    return generator
8
9num_shards = 4
10
11# Create a dataset of shard IDs
12shard_dataset = tf.data.Dataset.range(num_shards)
13
14# Interleave generators from each shard
15dataset = shard_dataset.interleave(
16    lambda shard_id: tf.data.Dataset.from_generator(
17        make_generator(shard_id.numpy(), num_shards),
18        output_signature=(
19            tf.TensorSpec(shape=(224, 224, 3), dtype=tf.float32),
20            tf.TensorSpec(shape=(), dtype=tf.int32),
21        )
22    ),
23    num_parallel_calls=tf.data.AUTOTUNE,
24    cycle_length=num_shards,
25    deterministic=False  # Allow out-of-order for better performance
26)
27
28dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)

interleave runs multiple generators simultaneously in separate threads.

Method 2: Generator for Paths, Map for Processing

Move heavy processing out of the generator into a parallelizable map function:

python

1def path_generator():
2    """Lightweight generator that yields file paths only."""
3    for i in range(10000):
4        yield f"image_{i}.jpg", i
5
6# Generator is fast — just yields paths
7path_dataset = tf.data.Dataset.from_generator(
8    path_generator,
9    output_signature=(
10        tf.TensorSpec(shape=(), dtype=tf.string),
11        tf.TensorSpec(shape=(), dtype=tf.int32),
12    )
13)
14
15def load_and_preprocess_tf(path, label):
16    """Heavy processing in a tf.py_function or pure TF ops."""
17    image = tf.io.read_file(path)
18    image = tf.image.decode_jpeg(image, channels=3)
19    image = tf.image.resize(image, [224, 224])
20    image = tf.cast(image, tf.float32) / 255.0
21    return image, label
22
23# Parallelize the heavy processing
24dataset = path_dataset.map(
25    load_and_preprocess_tf,
26    num_parallel_calls=tf.data.AUTOTUNE
27)
28
29dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)

The generator is now trivially fast, and all heavy I/O and preprocessing happens in the parallel map.

Method 3: tf.py_function in Map

If preprocessing requires Python libraries (OpenCV, PIL, custom code):

python

1import cv2
2
3def path_generator():
4    for i in range(10000):
5        yield f"image_{i}.jpg", i
6
7path_dataset = tf.data.Dataset.from_generator(
8    path_generator,
9    output_signature=(
10        tf.TensorSpec(shape=(), dtype=tf.string),
11        tf.TensorSpec(shape=(), dtype=tf.int32),
12    )
13)
14
15def preprocess_with_opencv(path, label):
16    def _process(path_bytes, label_val):
17        path_str = path_bytes.numpy().decode('utf-8')
18        img = cv2.imread(path_str)
19        img = cv2.resize(img, (224, 224))
20        img = img.astype(np.float32) / 255.0
21        return img, label_val.numpy()
22
23    image, label = tf.py_function(
24        _process,
25        [path, label],
26        [tf.float32, tf.int32]
27    )
28    image.set_shape((224, 224, 3))
29    label.set_shape(())
30    return image, label
31
32dataset = path_dataset.map(
33    preprocess_with_opencv,
34    num_parallel_calls=tf.data.AUTOTUNE
35)

tf.py_function wraps Python code to run in TensorFlow's data pipeline, bypassing the GIL for I/O-bound operations.

Method 4: Convert to TFRecords (Fastest)

For maximum performance, preprocess data once and store as TFRecords:

python

1# Write TFRecords
2writer = tf.io.TFRecordWriter('data.tfrecord')
3for i in range(10000):
4    image = load_and_preprocess(f"image_{i}.jpg")
5    label = get_label(i)
6
7    feature = {
8        'image': tf.train.Feature(float_list=tf.train.FloatList(value=image.flatten())),
9        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
10    }
11    example = tf.train.Example(features=tf.train.Features(feature=feature))
12    writer.write(example.SerializeToString())
13writer.close()
14
15# Read TFRecords with parallel I/O
16files = tf.data.Dataset.list_files('data_shard_*.tfrecord')
17dataset = files.interleave(
18    tf.data.TFRecordDataset,
19    num_parallel_calls=tf.data.AUTOTUNE,
20    cycle_length=4
21)
22
23def parse_example(serialized):
24    features = tf.io.parse_single_example(serialized, {
25        'image': tf.io.FixedLenFeature([224 * 224 * 3], tf.float32),
26        'label': tf.io.FixedLenFeature([1], tf.int64),
27    })
28    image = tf.reshape(features['image'], (224, 224, 3))
29    label = features['label'][0]
30    return image, label
31
32dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
33dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)

Complete Optimized Pipeline

python

1dataset = (
2    tf.data.Dataset.from_generator(path_generator, output_signature=...)
3    .shuffle(buffer_size=10000)
4    .map(load_and_preprocess_tf, num_parallel_calls=tf.data.AUTOTUNE)
5    .batch(32)
6    .prefetch(tf.data.AUTOTUNE)
7)
8
9model.fit(dataset, epochs=10)

Common Pitfalls

Putting heavy processing inside the generator: The generator runs in a single Python thread. Move I/O and preprocessing into .map() with num_parallel_calls=tf.data.AUTOTUNE so TensorFlow can parallelize it.
Forgetting prefetch(tf.data.AUTOTUNE): Without prefetch, the GPU idles while waiting for the next batch. prefetch overlaps data preparation with model training, keeping the GPU busy.
Using deterministic=True with interleave: Setting deterministic=True (the default) forces interleave to return elements in order, which serializes the generators and reduces parallelism. Set deterministic=False for training where order does not matter.
Not setting output shapes after tf.py_function: tf.py_function returns tensors with unknown shapes. Call tensor.set_shape(...) after the function to restore shape information, otherwise downstream operations (batching, model layers) may fail.
Generator not being re-entrant: TensorFlow may call the generator multiple times across epochs. If the generator uses external state (file handles, database connections), ensure it reinitializes properly. Use a factory function that returns a fresh generator each time.

Summary

from_generator is single-threaded — move heavy processing to .map() with num_parallel_calls
Use interleave with multiple generator instances for parallel data loading
Keep generators lightweight (yield paths/indices) and parallelize processing in map
Convert data to TFRecords for maximum I/O performance
Always use .prefetch(tf.data.AUTOTUNE) to overlap data loading and training

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

parallelising tf.data.Dataset.from_generator

ML System Design practice on Codemia

Introduction

The Problem

Method 1: Interleave Multiple Generators

Method 2: Generator for Paths, Map for Processing

Method 3: tf.py_function in Map

Method 4: Convert to TFRecords (Fastest)

Complete Optimized Pipeline

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia