TensorFlow
custom data loading
asynchronous computation
machine learning
deep learning

Tensorflow custom data load asynchronous computation

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

Introduction

TensorFlow model performance depends heavily on the input pipeline. If loading, decoding, or preprocessing is slow, the accelerator sits idle and training throughput drops even when the model itself is well optimized.

Use tf.data for custom loading

The standard way to build a custom pipeline in TensorFlow is tf.data. It lets you describe file reading, parsing, batching, and buffering as one pipeline that TensorFlow can optimize and run efficiently.

python
1import tensorflow as tf
2
3IMAGE_SIZE = (224, 224)
4
5def parse_example(path, label):
6    raw = tf.io.read_file(path)
7    image = tf.io.decode_jpeg(raw, channels=3)
8    image = tf.image.resize(image, IMAGE_SIZE)
9    image = tf.cast(image, tf.float32) / 255.0
10    return image, label
11
12paths = tf.constant(["data/cat.jpg", "data/dog.jpg"])
13labels = tf.constant([0, 1])
14
15dataset = tf.data.Dataset.from_tensor_slices((paths, labels))
16dataset = dataset.shuffle(1000)
17dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
18dataset = dataset.batch(32)
19dataset = dataset.prefetch(tf.data.AUTOTUNE)

This pipeline already does the important work. map can parse several samples in parallel, and prefetch overlaps data preparation with model execution so the next batch is ready sooner.

What asynchronous computation means in practice

Many TensorFlow questions ask whether asynchronous loading requires manual Python threads. Usually it does not. In most training jobs, the right approach is to describe the pipeline with tf.data operators and let TensorFlow schedule overlapping work.

With prefetch, the CPU can prepare the next batch while the GPU or TPU runs the current training step. That overlap is the practical meaning of asynchronous input computation for most users.

If your dataset spans many files, interleave helps keep I/O busy by opening more than one source at a time.

python
1import tensorflow as tf
2
3files = tf.data.Dataset.list_files("records/*.tfrecord", shuffle=True)
4dataset = files.interleave(
5    lambda path: tf.data.TFRecordDataset(path),
6    cycle_length=tf.data.AUTOTUNE,
7    num_parallel_calls=tf.data.AUTOTUNE,
8    deterministic=False,
9)
10dataset = dataset.prefetch(tf.data.AUTOTUNE)

That pattern is especially useful for TFRecord-based training jobs that would otherwise read files too sequentially.

When a generator is the right tool

Sometimes the source is not a normal file format. You may be reading from a simulator, a custom binary structure, or a streaming system. In those cases, Dataset.from_generator is a good escape hatch.

python
1import tensorflow as tf
2
3def sample_generator():
4    for i in range(10):
5        yield [float(i), float(i + 1)], i % 2
6
7dataset = tf.data.Dataset.from_generator(
8    sample_generator,
9    output_signature=(
10        tf.TensorSpec(shape=(2,), dtype=tf.float32),
11        tf.TensorSpec(shape=(), dtype=tf.int32),
12    ),
13)
14dataset = dataset.batch(4).prefetch(tf.data.AUTOTUNE)

This works well for unusual sources, but native TensorFlow readers are usually faster and easier to optimize than Python-driven generators.

Where to look next when throughput is still low

After the basic pipeline is in place, the next useful tools are cache, shuffle, and the TensorFlow Profiler. Caching can remove repeated decode work when the transformed dataset fits in memory or on local disk. Profiling shows whether time is being lost in file I/O, CPU preprocessing, or accelerator stalls. That is far more reliable than guessing based on how busy the training loop looks from Python.

Common Pitfalls

  • Writing your own Python threading layer instead of using tf.data primitives.
  • Forgetting prefetch, which prevents input work from overlapping with model execution.
  • Using from_generator for ordinary file loading when built-in readers would be faster.
  • Placing heavy Python logic inside map, which reduces optimization opportunities.
  • Tuning blindly without profiling whether the real bottleneck is I/O, CPU preprocessing, or accelerator utilization.

Summary

  • Use tf.data as the default way to build TensorFlow input pipelines.
  • 'map with parallel calls and prefetch are the core tools for asynchronous input work.'
  • 'interleave helps when records are spread across many files.'
  • 'from_generator is useful for unusual data sources, but built-in readers should be preferred when possible.'
  • Profile the pipeline before optimizing so you fix the real bottleneck.

Related reading
Free course
Beginner
7 lessons
2 hours
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

All Rights Reserved.