How to write and read TFRecord files in TensorFlow

Training data that needs to move between TensorFlow jobs is easier to reuse when each example is stored in a consistent record format. A TFRecord file stores serialized records that tf.data can stream back into an input pipeline without keeping the full dataset in memory.

The common TensorFlow pattern is to wrap each row in a tf.train.Example message, serialize it with SerializeToString(), and write it with tf.io.TFRecordWriter. Reading starts with tf.data.TFRecordDataset, which returns serialized byte records until a parser maps them back to typed feature tensors.

A small feature and label set keeps the file creation, raw record count, first parsed example, and first parsed batch visible in one run. Use the same schema names and types on the writer and reader sides; a key or dtype mismatch turns a valid TFRecord file into records that the parser cannot decode as intended.

1. Open a terminal in a Python environment where TensorFlow already imports cleanly.

   $ python3 -c "import tensorflow as tf; print(tf.__version__)"
   2.21.0

   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Save the TFRecord write/read demo as

   tfrecord_write_read.py

   .

   tfrecord_write_read.py

   import os
   from pathlib import Path
    
   os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
    
   import tensorflow as tf
    
    
   def float_feature(value):
       return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
    
    
   def int64_feature(value):
       return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
    
    
   rows = [
       {"feature_a": 0.10, "feature_b": 1.20, "label": 0},
       {"feature_a": 0.40, "feature_b": 0.70, "label": 1},
       {"feature_a": 0.30, "feature_b": 1.10, "label": 0},
       {"feature_a": 0.90, "feature_b": 0.20, "label": 1},
   ]
    
   record_path = Path("data/training.tfrecord")
   record_path.parent.mkdir(parents=True, exist_ok=True)
    
   with tf.io.TFRecordWriter(str(record_path)) as writer:
       for row in rows:
           example = tf.train.Example(
               features=tf.train.Features(
                   feature={
                       "feature_a": float_feature(row["feature_a"]),
                       "feature_b": float_feature(row["feature_b"]),
                       "label": int64_feature(row["label"]),
                   }
               )
           )
           writer.write(example.SerializeToString())
    
   feature_description = {
       "feature_a": tf.io.FixedLenFeature([], tf.float32),
       "feature_b": tf.io.FixedLenFeature([], tf.float32),
       "label": tf.io.FixedLenFeature([], tf.int64),
   }
    
    
   def parse_record(serialized_example):
       parsed = tf.io.parse_single_example(serialized_example, feature_description)
       label = parsed.pop("label")
       return parsed, label
    
    
   def read_records():
       return tf.data.TFRecordDataset(str(record_path))
    
    
   parsed_dataset = read_records().map(parse_record)
    
   record_count = sum(1 for _ in read_records())
   first_features, first_label = next(iter(parsed_dataset))
   batch_features, batch_labels = next(iter(parsed_dataset.batch(2)))
    
   first_record = {
       "feature_a": round(float(first_features["feature_a"].numpy()), 2),
       "feature_b": round(float(first_features["feature_b"].numpy()), 2),
       "label": int(first_label.numpy()),
   }
   first_batch_features = {
       name: [round(float(value), 2) for value in tensor.numpy().tolist()]
       for name, tensor in batch_features.items()
   }
    
   print(f"tfrecord_path={record_path}")
   print(f"file_exists={record_path.exists()}")
   print(f"record_count={record_count}")
   print(f"first_record={first_record}")
   print(f"first_batch_features={first_batch_features}")
   print(f"first_batch_labels={batch_labels.numpy().tolist()}")

   The TF_CPP_MIN_LOG_LEVEL line keeps TensorFlow startup messages out of this short data-format check. It does not change how the TFRecord file is written or parsed.

3. Check that the writer and parser use the same feature names and data types.

   feature_description = {
       "feature_a": tf.io.FixedLenFeature([], tf.float32),
       "feature_b": tf.io.FixedLenFeature([], tf.float32),
       "label": tf.io.FixedLenFeature([], tf.int64),
   }

   Change the reader schema whenever the written tf.train.Example fields change. A missing key, wrong dtype, or unexpected shape raises a parse error or feeds incorrect tensors into the model.

4. Run the script and confirm the file, record count, first parsed record, and first parsed batch.

   $ python3 tfrecord_write_read.py
   tfrecord_path=data/training.tfrecord
   file_exists=True
   record_count=4
   first_record={'feature_a': 0.1, 'feature_b': 1.2, 'label': 0}
   first_batch_features={'feature_a': [0.1, 0.4], 'feature_b': [1.2, 0.7]}
   first_batch_labels=[0, 1]

   Counting every record is useful for a small verification file. For large training sets, use the parsed dataset directly and rely on dataset metadata, job logs, or sampled checks instead of scanning the full file only to count rows.

5. Connect the parsed records to the training input stages after the round trip prints the expected fields.

   train_dataset = (
       tf.data.TFRecordDataset("data/training.tfrecord")
       .map(parse_record, num_parallel_calls=tf.data.AUTOTUNE)
       .shuffle(1000)
       .batch(32)
       .prefetch(tf.data.AUTOTUNE)
   )

   Put shuffle(), batch(), and prefetch(tf.data.AUTOTUNE) after parsing so the training code receives typed feature tensors and labels.
   Related: How to optimize TensorFlow data pipeline performance

6. Remove the demo script and sample record when they were created only for local verification.

   $ rm tfrecord_write_read.py data/training.tfrecord

   Keep the .tfrecord file instead when it is the training artifact that another job will consume.