TensorFlow
TFRecords
Machine Learning
Data Processing
Tutorials

TensorFlow - Read all examples from a TFRecords at once?

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

TensorFlow is designed to stream TFRecord examples through tf.data, not to load huge datasets into memory all at once. If you truly want every example from a TFRecord in one in-memory structure, you can do it, but it is usually best reserved for small files, debugging, or preprocessing steps where memory use is predictable.

The Normal Way Is to Stream

A TFRecord file contains serialized examples, and the usual pattern is:

  1. create a TFRecordDataset
  2. parse each record
  3. batch and prefetch
  4. feed the dataset into training or evaluation

That design keeps memory usage controlled and scales much better than trying to materialize the entire dataset in RAM.

Parse TFRecords into Tensors

Here is a small example that reads examples with two features:

python
1import tensorflow as tf
2
3
4feature_spec = {
5    "x": tf.io.FixedLenFeature([2], tf.float32),
6    "y": tf.io.FixedLenFeature([], tf.int64),
7}
8
9
10def parse_example(serialized):
11    example = tf.io.parse_single_example(serialized, feature_spec)
12    return example["x"], example["y"]
13
14
15dataset = tf.data.TFRecordDataset("data.tfrecord")
16dataset = dataset.map(parse_example)
17
18for x, y in dataset.take(3):
19    print(x.numpy(), y.numpy())

This is the scalable pattern. If the goal is model training, stop here and keep the pipeline streaming.

Materialize the Whole TFRecord When the File Is Small

If the file is small enough and you want all examples in Python memory, iterate over the dataset and collect the results:

python
1import tensorflow as tf
2
3
4feature_spec = {
5    "x": tf.io.FixedLenFeature([2], tf.float32),
6    "y": tf.io.FixedLenFeature([], tf.int64),
7}
8
9
10def parse_example(serialized):
11    example = tf.io.parse_single_example(serialized, feature_spec)
12    return example["x"], example["y"]
13
14
15dataset = tf.data.TFRecordDataset("data.tfrecord").map(parse_example)
16
17xs = []
18ys = []
19
20for x, y in dataset:
21    xs.append(x.numpy())
22    ys.append(y.numpy())
23
24print(len(xs), len(ys))

This reads every example and keeps it in Python lists. It is straightforward and fine for small data, but it defeats the main performance and memory benefits of TFRecords for large training sets.

Convert Everything to One Tensor Batch

If every example has the same fixed shape, you can also batch the whole dataset into one tensor after you know its size. For a small file, one way is to materialize the examples and stack them:

python
1import numpy as np
2
3X = np.stack(xs)
4y = np.array(ys)
5
6print(X.shape)
7print(y.shape)

If you already know the TFRecord contains exactly N examples, another option is:

python
all_in_one_batch = dataset.batch(1000)
for X_batch, y_batch in all_in_one_batch.take(1):
    print(X_batch.shape, y_batch.shape)

That only works cleanly when the dataset size is known and fits comfortably in memory.

Beware of Variable-Length Features

Reading “everything at once” gets harder if examples contain variable-length sequences, images of different shapes, or sparse features. In those cases, you may need:

  • padding
  • ragged tensors
  • sparse tensors
  • a custom collation step in Python

This is one reason the streaming tf.data approach is usually the better default. It lets you batch and transform records incrementally instead of solving one giant in-memory packing problem.

When It Makes Sense to Read Everything

Reading all examples at once is reasonable when:

  • the TFRecord file is tiny
  • you are writing a quick inspection script
  • you want to convert the data to another format once
  • the dataset must fit into another library that expects in-memory arrays

For normal model training, especially on real datasets, streaming is almost always the better architecture.

Common Pitfalls

The most common mistake is assuming TFRecords imply “load everything first.” They do not. They are meant to support efficient streaming pipelines.

Another pitfall is calling .numpy() on every element of a huge dataset and accidentally moving the bottleneck into Python and system memory.

It is also easy to forget that variable-length features complicate whole-dataset batching. Fixed-shape examples are much easier to collect into one tensor than ragged data.

Finally, do not optimize the wrong thing. If the use case is training, a well-batched tf.data pipeline is usually more efficient and more scalable than materializing the whole TFRecord into memory.

Summary

  • TFRecords are primarily designed for streaming input pipelines, not one-shot full-memory loading.
  • You can read all examples by iterating through TFRecordDataset and collecting parsed results.
  • Whole-dataset batching is practical only when the data is small and shape-consistent.
  • Variable-length features make “read everything at once” more complicated.
  • For training workloads, prefer the standard streaming tf.data pipeline.

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