Introduction

It is normal for much of the TensorFlow Dataset API to run on CPU instead of GPU. tf.data is primarily an input pipeline system for reading, parsing, decoding, shuffling, and batching data, and those operations often happen before tensors reach GPU compute kernels. The real question is not whether every dataset op uses GPU, but whether the overall training pipeline keeps the GPU fed efficiently.

Why tf.data Usually Runs on CPU

Most dataset work involves:

• file reading
• image decode
• parsing records
• string processing
• batching and prefetching

These tasks are typically CPU-oriented. GPU acceleration matters most for model forward and backward passes, not for every stage of input preparation.

So seeing low or zero GPU usage from the input pipeline itself is not automatically a problem.

Build a Healthy Input Pipeline

A practical baseline looks like this:

1import tensorflow as tf
2
3def preprocess(x, y):
4    x = tf.cast(x, tf.float32) / 255.0
5    return x, y
6
7(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
8
9ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
10ds = ds.shuffle(10_000)
11ds = ds.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
12ds = ds.batch(128)
13ds = ds.prefetch(tf.data.AUTOTUNE)

This does not force GPU execution for preprocessing, but it often improves end-to-end throughput because CPU and GPU can work concurrently.

GPU Usage Depends on the Training Step

The dataset itself is not the final consumer. The model training step is.

1model = tf.keras.Sequential([
2    tf.keras.layers.Flatten(input_shape=(28, 28)),
3    tf.keras.layers.Dense(128, activation="relu"),
4    tf.keras.layers.Dense(10)
5])
6
7model.compile(
8    optimizer="adam",
9    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
10    metrics=["accuracy"]
11)
12
13model.fit(ds, epochs=2)

If the model is small, the GPU may still look underutilized because compute per batch is tiny. That is a workload-size issue, not necessarily a dataset problem.

Check Whether TensorFlow Sees the GPU

First verify the runtime itself.

import tensorflow as tf

print(tf.config.list_physical_devices("GPU"))

If this returns an empty list, the issue is environment setup, not tf.data.

When Dataset Work Becomes the Bottleneck

tf.data can limit GPU throughput when preprocessing is too slow. Common causes:

• heavy Python logic inside map
• image decode or augmentation slower than training
• no prefetch
• batch size too small
• remote storage latency

A slow input pipeline makes the GPU wait idle between steps even if model ops themselves are correctly placed on GPU.

Prefer TensorFlow Ops Over Python Work

Avoid Python-heavy transformations inside the pipeline.

Less scalable pattern:

def bad_map(x, y):
    # Python or NumPy heavy logic here can slow the pipeline
    return x, y

Better pattern:

1def good_map(x, y):
2    x = tf.image.random_flip_left_right(x)
3    x = tf.image.resize(x, (128, 128))
4    return x, y

Using TensorFlow ops gives the runtime more room to optimize execution.

Useful Performance Features

Several tf.data settings help GPU utilization indirectly:

• 'num_parallel_calls=tf.data.AUTOTUNE'
• 'prefetch(tf.data.AUTOTUNE)'
• caching small datasets with .cache()
• batching before expensive operations when appropriate

Example with cache:

1ds = ds.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
2ds = ds.cache()
3ds = ds.batch(128)
4ds = ds.prefetch(tf.data.AUTOTUNE)

Cache helps when the dataset fits in memory and preprocessing is deterministic enough to reuse.

Profiling Instead of Guessing

If GPU usage looks poor, profile the pipeline rather than assuming dataset ops should move to GPU.

Useful questions:

• Is GPU visible to TensorFlow
• Is step time dominated by input or by model compute
• Does throughput improve with larger batch size
• Does prefetch reduce idle time

The fix is usually pipeline tuning, not forcing dataset ops onto GPU.

Common Pitfalls

• Treating CPU execution of dataset ops as automatic misconfiguration.
• Forgetting prefetch and leaving GPU starved between steps.
• Using Python or NumPy inside map and slowing the pipeline.
• Testing with tiny models or tiny batches and expecting high GPU utilization.
• Diagnosing input performance before confirming TensorFlow can even see the GPU.

Summary

• Most tf.data work runs on CPU by design, and that is normal.
• The real goal is keeping the GPU busy during model execution.
• Use parallel mapping, batching, caching, and prefetching to improve throughput.
• Verify GPU visibility separately from dataset behavior.
• Profile end-to-end training performance before deciding the pipeline is the problem.