TensorFlow
TFRecord
tf.dataset
data processing
machine learning

Write tf.dataset back to TFRecord

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

Writing a tf.data.Dataset back to TFRecord is a two-step job: convert each dataset element into a serialized record, then stream those serialized bytes to a writer. The important part is that TFRecord stores bytes, not arbitrary Python objects, so you need to define an explicit record format before writing anything.

Think in Terms of Serialized Examples

TFRecord is just a container format for binary records. In TensorFlow workflows, those records are usually serialized tf.train.Example messages.

That means a dataset like:

  • image tensor
  • label integer
  • id string

must be turned into a byte string for each element before it can be written.

Create a Serializer for One Record

The most common pattern is to write a Python function that converts one example into tf.train.Example, then serialize it.

python
1import tensorflow as tf
2
3
4def _bytes_feature(value: bytes) -> tf.train.Feature:
5    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
6
7
8def _int64_feature(value: int) -> tf.train.Feature:
9    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
10
11
12def serialize_example(item_id: bytes, label: int) -> bytes:
13    example = tf.train.Example(
14        features=tf.train.Features(
15            feature={
16                "item_id": _bytes_feature(item_id),
17                "label": _int64_feature(label),
18            }
19        )
20    )
21    return example.SerializeToString()

This function is pure serialization logic. Keep it narrow and predictable.

Write a Dataset to TFRecord

Once you can serialize one example, you can iterate through a dataset and write the records.

python
1import tensorflow as tf
2
3item_ids = tf.constant([b"a100", b"a101", b"a102"])
4labels = tf.constant([1, 0, 1], dtype=tf.int64)
5
6dataset = tf.data.Dataset.from_tensor_slices((item_ids, labels))
7
8path = "/tmp/items.tfrecord"
9with tf.io.TFRecordWriter(path) as writer:
10    for item_id, label in dataset:
11        record = serialize_example(
12            item_id=item_id.numpy(),
13            label=int(label.numpy()),
14        )
15        writer.write(record)
16
17print("wrote", path)

This version is simple and easy to debug. It works well when dataset elements are manageable in eager execution.

TensorFlow-Style Pipeline Writing

If you want to stay closer to tf.data, map each record to a serialized scalar string and then pass that dataset into a writer.

python
1import tensorflow as tf
2
3
4def tf_serialize_example(item_id, label):
5    serialized = tf.py_function(
6        func=lambda x, y: serialize_example(x.numpy(), int(y.numpy())),
7        inp=[item_id, label],
8        Tout=tf.string,
9    )
10    serialized.set_shape(())
11    return serialized
12
13
14serialized_dataset = dataset.map(tf_serialize_example)
15writer = tf.data.experimental.TFRecordWriter("/tmp/items_from_dataset.tfrecord")
16writer.write(serialized_dataset)

The explicit loop is easier to reason about. The dataset-based writer is convenient when you want to keep the output stage inside a TensorFlow input pipeline.

Include Tensor Data Safely

If dataset elements contain tensors such as images or embeddings, serialize them explicitly instead of trying to flatten them by hand.

python
1import numpy as np
2import tensorflow as tf
3
4
5def serialize_tensor_example(vector: tf.Tensor, label: int) -> bytes:
6    example = tf.train.Example(
7        features=tf.train.Features(
8            feature={
9                "vector": _bytes_feature(tf.io.serialize_tensor(vector).numpy()),
10                "label": _int64_feature(label),
11            }
12        )
13    )
14    return example.SerializeToString()
15
16
17vector = tf.constant(np.array([0.1, 0.2, 0.3], dtype=np.float32))
18record = serialize_tensor_example(vector, 1)
19print(len(record))

This preserves dtype and shape information in a form TensorFlow can reconstruct later with tf.io.parse_tensor.

Always Verify by Reading the File Back

Do not assume the writer is correct just because the file exists. Read the TFRecord back and inspect a few examples.

python
1import tensorflow as tf
2
3feature_spec = {
4    "item_id": tf.io.FixedLenFeature([], tf.string),
5    "label": tf.io.FixedLenFeature([], tf.int64),
6}
7
8
9def parse_record(record):
10    parsed = tf.io.parse_single_example(record, feature_spec)
11    return parsed["item_id"], parsed["label"]
12
13
14raw_dataset = tf.data.TFRecordDataset("/tmp/items.tfrecord")
15for item_id, label in raw_dataset.map(parse_record):
16    print(item_id.numpy(), int(label.numpy()))

Round-trip validation catches schema mistakes immediately.

Common Pitfalls

The most common mistake is trying to write dataset elements directly without serializing them first. TFRecord stores bytes, so each element must become a serialized record.

Another issue is mixing Python objects and tensors without a clear schema. If you cannot describe the record format precisely, you will have trouble parsing it later.

A third problem is skipping verification. A TFRecord file can be successfully written and still be unusable because the feature names, dtypes, or shapes are inconsistent.

Summary

  • Convert each dataset element into a serialized byte record before writing.
  • 'tf.train.Example is the standard format for most TFRecord workflows.'
  • Use a simple eager loop first if you want the easiest implementation to debug.
  • Serialize tensors explicitly with tf.io.serialize_tensor when needed.
  • Always read a sample of the file back to confirm the schema is correct.

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.