Introduction

TFRecord can help Keras training, but the speedup comes from the whole tf.data pipeline rather than from the file format by itself. To make training fast, you need parallel reading, efficient parsing, batching, and prefetching so input work overlaps with model execution.

Build a Correct Baseline First

Before optimizing, make sure the parser is correct and the batch shapes are what the model expects.

1import tensorflow as tf
2
3feature_spec = {
4    "x": tf.io.FixedLenFeature([32], tf.float32),
5    "y": tf.io.FixedLenFeature([], tf.int64),
6}
7
8
9def parse_example(record):
10    ex = tf.io.parse_single_example(record, feature_spec)
11    return ex["x"], ex["y"]
12
13files = tf.io.gfile.glob("data/train-*.tfrecord")
14ds = tf.data.TFRecordDataset(files)
15ds = ds.map(parse_example)
16ds = ds.batch(128)
17
18for batch_x, batch_y in ds.take(1):
19    print(batch_x.shape, batch_y.shape)

Once this is correct, performance tuning becomes meaningful.

Add Parallel Reads, Parallel Map, and Prefetch

The biggest wins usually come from overlapping I/O and parsing with model execution.

1import tensorflow as tf
2
3files = tf.io.gfile.glob("data/train-*.tfrecord")
4
5ds = tf.data.Dataset.from_tensor_slices(files)
6ds = ds.shuffle(len(files))
7ds = ds.interleave(
8    lambda f: tf.data.TFRecordDataset(f),
9    cycle_length=tf.data.AUTOTUNE,
10    num_parallel_calls=tf.data.AUTOTUNE,
11)
12ds = ds.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
13ds = ds.shuffle(20000)
14ds = ds.batch(256, drop_remainder=True)
15ds = ds.prefetch(tf.data.AUTOTUNE)

prefetch lets the input pipeline prepare future batches while the accelerator is busy on the current step. AUTOTUNE lets TensorFlow pick a reasonable level of parallelism.

Shard TFRecords Well

How the TFRecords were written matters too. A good dataset layout usually has:

• multiple shard files rather than one giant file
• examples that are cheap to parse
• consistent record structure
• only the features that training actually needs

If the records are badly sharded or overly expensive to parse, the input pipeline stays slow even with correct tf.data operators.

Use Cache Carefully

Caching can help when the dataset fits in memory or when preprocessing is expensive and repeated every epoch.

cached_ds = ds.cache().prefetch(tf.data.AUTOTUNE)

But cache() is not a free win. If the dataset is too large, it can create memory pressure and make the system less stable. Use it only after measuring the memory tradeoff.

Measure the Input Pipeline Directly

Do not assume the pipeline is fast because the code looks optimized. Measure samples per second outside the model.

1import time
2
3start = time.time()
4count = 0
5
6for x, y in ds.take(200):
7    count += x.shape[0]
8
9elapsed = time.time() - start
10print("samples/sec:", count / elapsed)

If the GPU or TPU is idle between steps, the dataset is still not feeding the model efficiently enough.

Feed the Dataset Directly Into Keras

Once the pipeline is tuned, connect it directly to model.fit.

1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(32,)),
3    tf.keras.layers.Dense(64, 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=5)

This avoids slow Python-side generators and lets TensorFlow optimize the dataset path end to end.

Avoid Python in the Hot Path

A frequent hidden slowdown is leaving expensive Python work in the input pipeline or falling back to Python generators around model.fit. The more the input path stays inside tf.data transformations, the more TensorFlow can overlap and optimize execution.

That is why moving parsing, batching, and prefetching into the dataset pipeline usually matters more than small model-side tweaks.

Common Pitfalls

The biggest mistake is blaming TFRecord when the real bottleneck is a weak tf.data pipeline.

Another issue is forgetting prefetch, which leaves the accelerator waiting for the next batch.

A third problem is using cache() blindly on datasets that do not fit comfortably in memory.

Summary

• TFRecord speed depends on the full input pipeline, not only on file format.
• Start from a correct parser and correct batch shapes.
• Use interleave, parallel map, and prefetch(tf.data.AUTOTUNE) to overlap work.
• Shard TFRecords well so the input pipeline can read efficiently.
• Measure throughput directly instead of assuming the pipeline is fast.