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
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:
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:
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):
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:
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
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
