TFRecords & Dataset API

Training throughput is usually set by input, not compute. TFRecords store examples in a compact binary format that streams efficiently, and tf.data.Dataset is the pipeline that reads, decodes, augments, and batches them. Learned together, they keep the accelerator fed.

The idea

• A TFRecord file is a sequence of length-prefixed records; here each record is a serialized tf.train.Example.
• An Example maps string keys to tf.train.Feature values, which hold a BytesList, FloatList, or Int64List. Images are stored as raw encoded bytes, not decoded arrays.
• The feature schema lives in your code, not the file, so reads need tf.io.FixedLenFeature([], tf.string) and friends to parse.
• tf.data is lazy and graph-based: map, batch, and prefetch build a pipeline the runtime can parallelize.

Ordering rules of thumb: shuffle before batch, cache before an expensive map, prefetch(tf.data.AUTOTUNE) last, and num_parallel_calls on any map doing real work.

Worked example

Write 1,000 tiny (28, 28) float arrays.

• Each example holds {"image": bytes_feature(x.tobytes()), "label": int64_feature(y)}.
• Write in shards of about 200 records: part-00000.tfrecord, and so on.
• Read with tf.data.TFRecordDataset(files), then map(_parse), shuffle(1000), batch(32), prefetch(AUTOTUNE).
• _parse sets image = tf.io.decode_raw(f["image"], tf.float32) and reshapes to (28, 28).

Sharding matters: one giant file serializes reads, while many files let interleave run in parallel.

In code

import numpy as np, tensorflow as tf

def _bytes(x): return tf.train.Feature(bytes_list=tf.train.BytesList(value=[x]))
def _int(x):   return tf.train.Feature(int64_list=tf.train.Int64List(value=[x]))

with tf.io.TFRecordWriter("/tmp/data/part-00000.tfrecord") as w:
    for i in range(1000):
        img = np.random.rand(28, 28).astype("float32")
        ex = tf.train.Example(features=tf.train.Features(feature={
            "image": _bytes(img.tobytes()),
            "label": _int(i % 10),
        }))
        w.write(ex.SerializeToString())

def parse(proto):
    f = tf.io.parse_single_example(proto, {
        "image": tf.io.FixedLenFeature([], tf.string),
        "label": tf.io.FixedLenFeature([], tf.int64),
    })
    return tf.reshape(tf.io.decode_raw(f["image"], tf.float32), (28, 28)), f["label"]

ds = (tf.data.TFRecordDataset(["/tmp/data/part-00000.tfrecord"])
        .map(parse, num_parallel_calls=tf.data.AUTOTUNE)
        .shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE))

print(next(iter(ds))[0].shape)   # (32, 28, 28)

Check yourself

1. Why does the reader need the feature schema when the writer already serialized the data?
2. Why should shuffle come before batch?
3. What does prefetch(AUTOTUNE) overlap, and why does it help?

Key takeaways

• TFRecords store raw bytes; the parse schema reconstructs structure at read time.
• tf.data builds a lazy pipeline that overlaps input work with training steps.
• Sharding, shuffle-before-batch, and prefetch are the three highest-value defaults.