TensorFlow
TFRecords
Data Processing
Record Counting
Machine Learning

Obtaining total number of records from .tfrecords file in Tensorflow

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

TFRecord files do not store a built-in row count that you can read instantly like a database table statistic. If you need the total number of records, the normal approach is to iterate through the file and count them. That sounds inefficient, but it is the correct and dependable method unless you store the count separately when the file is created.

Count records with TFRecordDataset

In modern TensorFlow, the most direct method is to build a TFRecordDataset and count each element.

python
1import tensorflow as tf
2
3dataset = tf.data.TFRecordDataset("data.tfrecord")
4count = sum(1 for _ in dataset)
5
6print(count)

This works because each element yielded by TFRecordDataset corresponds to one serialized record from the file. It is simple, readable, and works for both small and large files, though it still has to scan the data.

Use reduce to keep the counting inside TensorFlow

If you prefer a TensorFlow-native counting step instead of a Python generator loop, use Dataset.reduce.

python
1import tensorflow as tf
2
3dataset = tf.data.TFRecordDataset("data.tfrecord")
4count = dataset.reduce(
5    tf.constant(0, dtype=tf.int64),
6    lambda acc, _: acc + 1
7)
8
9print(int(count.numpy()))

This does the same logical work but keeps the accumulation in the dataset pipeline. It can be a cleaner fit if the rest of your input pipeline is already expressed in TensorFlow operations.

Count across multiple TFRecord files

If the dataset is split across many files, pass a list of file names.

python
1import tensorflow as tf
2
3files = ["part-000.tfrecord", "part-001.tfrecord", "part-002.tfrecord"]
4dataset = tf.data.TFRecordDataset(files)
5count = sum(1 for _ in dataset)
6
7print(count)

This gives you the total number of records across all files, not just per file. That is often what you want when computing steps per epoch or validating that a data export finished correctly.

Compression settings matter

If the TFRecord file was written with compression, you need to provide the matching compression type while reading.

python
1import tensorflow as tf
2
3dataset = tf.data.TFRecordDataset(
4    "compressed.tfrecord",
5    compression_type="GZIP"
6)
7
8count = sum(1 for _ in dataset)
9print(count)

If the compression type is wrong, the read may fail or return invalid data. Counting records is still just iteration, but the dataset must be opened correctly.

When counting becomes expensive

Counting requires a full scan of the file, so it can take time on large datasets. If you need the count frequently, a better design is to save metadata when you write the records.

For example, when generating TFRecords, you might also write a JSON sidecar file:

python
1import json
2
3metadata = {"record_count": 125000}
4
5with open("data.tfrecord.meta.json", "w", encoding="utf-8") as f:
6    json.dump(metadata, f)

Then later you can read the stored count instantly instead of rescanning the binary file every time.

Common Pitfalls

The biggest mistake is assuming a TFRecord file has a quick built-in count query. It does not. If you did not store the count elsewhere, you have to iterate.

Another issue is forgetting compression settings. A GZIP-compressed TFRecord cannot be read correctly with the default dataset settings.

Developers also confuse "number of records" with "number of examples after parsing and filtering." If your pipeline later filters invalid records, the final usable example count can be smaller than the raw record count.

Finally, do not repeatedly rescan huge TFRecord datasets in tight loops or startup paths if you can avoid it. Persist the metadata once and reuse it.

Summary

  • TFRecord files do not expose a cheap built-in row-count lookup.
  • Use TFRecordDataset and count elements to get the record total.
  • 'Dataset.reduce is a clean TensorFlow-native way to do the same thing.'
  • Match the compression_type when reading compressed TFRecords.
  • If you need the count often, store it separately when writing the dataset.

Course illustration
Course illustration

All Rights Reserved.