1. Introduction to TFRecord
Nowadays, most Machine Learning (ML) and Deep Learning (DL) projects deal with large datasets. For example, GPT-3 was trained on 570GB data. Storing and reading large amounts of data efficiently is necessary in large-scale AI projects. TensorFlow provides a simple, efficient and flexible binary format called TFRecord. It contains a sequence of binary records of varying sizes, the length of the record, a checksum checking if the length was not corrupted and the actual data. Unlike text formats like CSV or JSON, TFRecord leverages Protocol Buffers (protobufs) to serialize data, making it more compact and faster to read and write.
In TFRecord, data is stored as serialized tf.train.Example objects using Protocol Buffers. This makes the data more portable and ensures compatibility with TensorFlow’s input pipeline. Whether datatype you are working with structured or unstructured data. TFRecord can efficiently store your dataset for optimal performance.
2. Why Use TFRecord?
The TFRecord format is specially designed to handle large-scale, complex datasets, TFRecord holds several advantages over other text formats:
1. Efficient Storage
As the TFRecord serializes the using binary leveraging Protocol Buffers, it uses less storage space than regular formats such as CSV or JSON. This helps reduce disk storage space, allowing more efficient use of storage resources.
2. Faster I/O Operations
When training machine learning models, reading data from disk can become a bottleneck, especially with large datasets. Binary formats like TFRecord are faster to read and write compared to text-based formats, as they are more straightforward for the CPU to process.
3. Flexibility in Storing Complex Data
TFRecord can handle various data types, including structured or unstructured data. You can store different types of data in a single .tfrecord file, for example, you can store text and image data in the same file, making it very useful for special use cases.
4. Great Integration with TensorFlow’s tf.data API
The TFRecord format is specifically optimized for use with TensorFlow’s tf.data API. It supports all the input pipeline features such as shuffling, prefetching, batching, and parallel data loading, ensuring efficient data pipeline from the data loading to the AI model.
3. Using TFRecord in TensorFlow
1. Writing Data to TFRecord
To use TFRecord, you first need to serialize your data into tf.train.Example messages. Here’s how to do it:
import tensorflow as tf
data = [
(1, b'first record'),
(2, b'second record'),
(3, b'third record')
]
def Serializing(feature0, feature1):
feature = {
'feature0': tf.train.Feature(int64_list=tf.train.Int64List(value=[feature0])),
'feature1': tf.train.Feature(bytes_list=tf.train.BytesList(value=[feature1]))
}
example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
return example_proto.SerializeToString()
# Writing data to a TFRecord file
tfrecord_filename = 'my_data.tfrecord'
with tf.io.TFRecordWriter(tfrecord_filename) as writer:
for feature0, feature1 in data:
example = serializing(feature0, feature1)
writer.write(example)
The Serializing() function converts each data entry into a tf.train.Example protocol buffer. The tf.io.TFRecordWriter then writes each serialized example to a .tfrecord file.
2. Reading and Parsing TFRecord Files
After writing data to TFRecord, you can read it back using the tf.data.TFRecordDataset and parse it with a feature description.
feature_description = {
'feature0': tf.io.FixedLenFeature([], tf.int64),
'feature1': tf.io.FixedLenFeature([], tf.string)
}
# Parsing function
def Parsing(example_proto):
return tf.io.parse_single_example(example_proto, feature_description)
dataset = tf.data.TFRecordDataset(tfrecord_filename)
dataset = dataset.map(parse_example)
for record in dataset:
print(record)
This code defines the structure of the serialized data using feature_description. The Parsing() function decodes each record, and the dataset is read using tf.data.TFRecordDataset.
3. Enhancing Input Pipelines with tf.data
As mentioned earlier the real power of TFRecord comes in its integration with the tf.data API. You can add any transformations such as shuffling, batching, and prefetching to optimize the input pipeline:
dataset = tf.data.TFRecordDataset(tfrecord_filename)
dataset = dataset.map(parse_example)
dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)
This ensures that data is efficiently loaded and processed in parallel with model training.
4. When Should You Use TFRecord?
While TFRecord provides many advantages, it is most beneficial in specific scenarios:
Large Datasets: When working with datasets too large to fit into memory, TFRecord’s binary format offers efficient storage and access.
Image and Audio Data: If you are dealing with complex data types like images or audio files, storing them in TFRecord format avoids the overhead of repeated file I/O operations.
Distributed Training: When using multi-GPU setups, TFRecord helps create efficient and high-speed data pipelines that can keep up with the demands of distributed training.
However, for smaller datasets or when AI project is straightforward, using regular text formats like CSV or JSON may be sufficient.
5. Summary
The TFRecord format is a powerful tool in the TensorFlow project for handling large, complex datasets efficiently. By using binary format, it offers faster read and write from disks, supports various data types, and integrates perfectly with TensorFlow’s tf.data API to build optimized input pipelines. While not necessary for every project, TFRecord becomes essential when dealing with large datasets and complex data types. If you have not yet, I suggest you use TFRecord in your next large-scale TensorFlow project to take advantage of its benefits.
References
Géron, A. (2019). Hands-On Machine Learning with Scikit-Learn, Keras & TensorFlow (2nd ed.). O’Reilly Media.