Introduction

In TensorFlow, prefetching is not something you add inside a Python function itself. It is a property of the tf.data pipeline around that function. The usual pattern is to map your custom loader or transform across a dataset, batch if needed, and then call .prefetch(...) so input preparation overlaps with model execution.

The Pipeline Structure

A typical performant pipeline looks like this:

1. create a Dataset
2. map a parsing or loading function across elements
3. batch the results
4. prefetch future batches

The important point is that prefetching happens on the dataset, not inside your file-reading or preprocessing function.

Using a Custom Python Function

If your custom loader cannot be expressed with pure TensorFlow ops, you can wrap it with tf.py_function. That lets Python run inside the input pipeline, but you should understand the tradeoff: Python functions are harder for TensorFlow to optimize and can become the bottleneck if they do heavy work.

Here is a runnable example that loads small NumPy files from disk with a Python function and then prefetches batches.

1import os
2import tempfile
3import numpy as np
4import tensorflow as tf
5
6# Create sample files for the demo.
7tmpdir = tempfile.mkdtemp()
8paths = []
9for i in range(6):
10    path = os.path.join(tmpdir, f"sample_{i}.npy")
11    np.save(path, np.array([i, i + 1, i + 2], dtype=np.float32))
12    paths.append(path)
13
14
15def load_numpy_file(path_tensor):
16    path = path_tensor.numpy().decode("utf-8")
17    return np.load(path).astype(np.float32)
18
19
20def tf_load_numpy(path_tensor):
21    tensor = tf.py_function(load_numpy_file, [path_tensor], Tout=tf.float32)
22    tensor.set_shape([3])
23    return tensor
24
25
26dataset = tf.data.Dataset.from_tensor_slices(paths)
27dataset = dataset.map(tf_load_numpy, num_parallel_calls=tf.data.AUTOTUNE)
28dataset = dataset.batch(2)
29dataset = dataset.prefetch(tf.data.AUTOTUNE)
30
31for batch in dataset.take(2):
32    print(batch)

This does two separate performance-related things:

• 'num_parallel_calls=tf.data.AUTOTUNE lets mapping happen with parallelism'
• '.prefetch(tf.data.AUTOTUNE) overlaps later input preparation with model work'

Why Prefetch Helps

Without prefetching, training tends to alternate between two phases:

• wait for the next batch to be prepared
• run the model on that batch

With prefetching, TensorFlow can prepare a future batch while the current one is already being consumed by the model. That reduces idle time, especially when input reading is slower than the actual forward or backward pass.

When tf.py_function Is a Problem

tf.py_function is useful, but it is not the ideal long-term path for every pipeline. Because it executes Python code, it can limit portability and throughput.

If possible, prefer TensorFlow-native ops such as:

• 'tf.io.read_file'
• 'tf.image.decode_png'
• 'tf.io.parse_tensor'
• 'Dataset.interleave for reading many files efficiently'

For example, if your data is text or images, pure TensorFlow parsing often performs better than jumping back to Python for every element.

A Cleaner Native Pattern When Possible

If the work can stay in TensorFlow, the pipeline is simpler and easier to optimize:

1dataset = tf.data.Dataset.list_files("/data/*.jpg")
2dataset = dataset.map(load_and_decode_image, num_parallel_calls=tf.data.AUTOTUNE)
3dataset = dataset.batch(32)
4dataset = dataset.prefetch(tf.data.AUTOTUNE)

This is usually where you want to end up if performance matters.

Common Pitfalls

The most common mistake is expecting .prefetch(...) to speed up a pipeline whose custom Python function is still the dominant bottleneck. Prefetch hides waiting, but it does not make expensive Python code free.

Another mistake is forgetting to set the shape after tf.py_function. TensorFlow often loses static shape information there, and downstream layers may fail or infer shapes poorly.

A third pitfall is calling .prefetch(...) too early in the pipeline and then assuming batching or expensive mapping is fully overlapped. Prefetch is most useful near the end of the input pipeline, typically after batching.

Summary

• Prefetching is configured on the tf.data.Dataset, not inside the custom Python function.
• Wrap truly Python-only logic with tf.py_function, then map, batch, and prefetch.
• Use num_parallel_calls=tf.data.AUTOTUNE together with .prefetch(tf.data.AUTOTUNE) for a strong default.
• Prefer TensorFlow-native ops when possible because they optimize better than Python callbacks.
• After tf.py_function, set the tensor shape explicitly so the rest of the pipeline stays well defined.