How to write and read TFRecord files in TensorFlow

Machine-learning input pipelines need a storage boundary that preserves feature values between data preparation and training. TFRecord provides that boundary as a sequential record file that TensorFlow can stream through tf.data without loading an entire dataset at once.

Each record may contain any byte string, but tf.train.Example provides a standard mapping from feature names to byte, float, or integer lists. tf.io.TFRecordWriter stores the serialized examples, while tf.data.TFRecordDataset returns those same examples as scalar string tensors for parsing.

The writer and reader must share the same feature names, shapes, and data types. The small weather dataset below uses two float inputs and one integer label, then reads all three rows into batches so a schema mismatch or truncated round trip fails before the data reaches a model.

1. Create tfrecord_round_trip.py with the sample rows and feature encoders.

   tfrecord_round_trip.py

   import os
   from pathlib import Path
    
   os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
    
   import tensorflow as tf
    
    
   record_path = Path("weather.tfrecord")
   rows = [
       {"temperature_c": 18.5, "rainfall_mm": 3.25, "will_rain": 1},
       {"temperature_c": 21.0, "rainfall_mm": 0.0, "will_rain": 0},
       {"temperature_c": 16.75, "rainfall_mm": 8.5, "will_rain": 1},
   ]
    
    
   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]))

   tf.train.Example accepts feature values through BytesList, FloatList, or Int64List. String and encoded-file schemas use a corresponding byte-feature helper.

2. Append the TFRecord serialization section to tfrecord_round_trip.py.

   def serialize_row(row):
       features = {
           "temperature_c": float_feature(row["temperature_c"]),
           "rainfall_mm": float_feature(row["rainfall_mm"]),
           "will_rain": int64_feature(row["will_rain"]),
       }
       example = tf.train.Example(
           features=tf.train.Features(feature=features)
       )
       return example.SerializeToString()
    
    
   with tf.io.TFRecordWriter(str(record_path)) as writer:
       for row in rows:
           writer.write(serialize_row(row))

   Opening weather.tfrecord replaces any existing file at that path. A separate filename preserves an earlier dataset.

3. Append the matching feature specification and parser to tfrecord_round_trip.py.

   feature_spec = {
       "temperature_c": tf.io.FixedLenFeature([], tf.float32),
       "rainfall_mm": tf.io.FixedLenFeature([], tf.float32),
       "will_rain": tf.io.FixedLenFeature([], tf.int64),
   }
    
    
   def parse_record(serialized_record):
       parsed = tf.io.parse_single_example(serialized_record, feature_spec)
       label = parsed.pop("will_rain")
       return parsed, label

   Each FixedLenFeature name and dtype must align with the serialized feature. An absent key, incompatible dtype, or wrong shape raises a parsing error instead of producing the intended tensors.

4. Append the batched TFRecord reader to tfrecord_round_trip.py.

   dataset = (
       tf.data.TFRecordDataset(str(record_path))
       .map(parse_record)
       .batch(2)
   )
    
   for batch_number, (features, labels) in enumerate(dataset, start=1):
       print(
           f"batch={batch_number} "
           f"temperature_c={features['temperature_c'].numpy().tolist()} "
           f"rainfall_mm={features['rainfall_mm'].numpy().tolist()} "
           f"will_rain={labels.numpy().tolist()}"
       )

5. Run the completed TFRecord round-trip program to confirm every row is parsed.

   $ python3 tfrecord_round_trip.py
   batch=1 temperature_c=[18.5, 21.0] rainfall_mm=[3.25, 0.0] will_rain=[1, 0]
   batch=2 temperature_c=[16.75] rainfall_mm=[8.5] will_rain=[1]

   Both batches come from parsing weather.tfrecord rather than the original Python list. Large training datasets should normally use multiple TFRecord shards so readers can parallelize file access.