Data Formats for Training in TensorFlow: Parquet, Petastorm, Feather, and More | by Chaim Rand | Towards Data Science

Machine learning is all about the data. To successfully train a sophisticated model you will need a high quality training dataset; a dataset that is sufficiently large, accurately labeled, and correctly represents the distribution of data samples in the real world. However, no less important is proper management of the data. By data management we are referring to how and where the data is stored, the ways in which it is accessed, and the transformations it undergoes during the development life-cycle. The focus of this post is on the file format used to store the training data and the implications it can have on the model training. Choosing the file format is one of the many important decisions you will need to make when defining your machine learning project.
This post is comprised of four sections. In the first section we will identify some of the properties that we would like our file format to have. In the second section we will review potential file formats and evaluate them against the desired properties we have found. Next we will survey different options for streaming these formats into a TensorFlow training session. In the last section we will describe a few experiments we ran to test some of these options.
Before we dive in, let’s set the stage. We will assume in this post that the size of our training dataset necessitates storing it in a distributed storage system, e.g. hundreds of terabytes. In the examples below we will use Amazon S3 for data storage, but the principals applies equally to any other distributed storage system. The data will likely be accessed by several consumers during the development cycle and for various purposes including data creation, data analysis, model training, and more. Our focus will be on training in TensorFlow, although much of what we will say pertains to other training frameworks as well.
The underlying assumption of this post is that a file format for our data is required, i.e. that we require one or more files for: 1. grouping together all elements/columns/features of single data samples, as well as 2. grouping together multiple data samples. Of course, one could envision scenarios in which one might choose to maintain all elements of all samples in individual files and in their raw format. However, in many cases, especially if the file sizes are small (e.g. several KBs) this strategy is likely to severely degrade the runtime performance of the data access.
The intention of this post is to highlight some of the important considerations when choosing a file format and to discuss a few of the different options available today. We will mention a number of formats as well as a number of software development frameworks and tools. These mentions should not be interpreted as an endorsement. The appropriate choice for you is likely to be based on a wide range of considerations some of which may be beyond the scope of this discussion. The post will include a number of simple code snippets. These are presented for the purposes of demonstration only and should not be viewed as optimal implementations.
The landscape of machine learning tools is extremely dynamic. Some of the formats and tools we will mention continue to evolve and some of the comments we will make may become outdated by the time you read this. Take care to keep track of announcements of new versions and new tools and be sure to make your design decisions on the most up to date information available.
Please do not hesitate to reach out to me with any comments or corrections.
To facilitate this discussion let’s take a look at a common file format used for training in TensorFlow, the TFRecord format.
TFRecord is a format based on protocol buffers specifically designed for use with TensorFlow. A TFRecord file is comprised of sequences of serialized binary samples. Each sample represents a tf.train.Example which, in turn, represents a dictionary of string to value mappings. The sequential nature of the TFRecord format facilitates high data streaming throughput. In particular, one does not need to download and open a full file in order to start traversing its contents. Moreover, the TensorFlow tf.data module includes the highly optimized TFRecordDataset class for building input pipelines based on data stored in TFRecord files. However, the TFRecord format is not without its faults. Here we name a few:
Single purpose: The TFRecord format is unlikely to satisfy the needs of any of the other data consumers on the development pipeline. It is common for development teams to maintain their data in a different format and create a derivative of the data in TFRecord format specifically for the purpose of training. Storing multiple copies of your data is not ideal, especially when your data set is large. Aside from the added cost of data storage, the implication is that every time there is a change to the core data set a new TFRecord copy needs to be generated. We may wish update the core data records based on findings from the training session, in which case we will also need to maintain a mapping between the generated TFRecord records and their corresponding entries in the core data format.
Extracting partial record data: Another problem with the TFRecord format arises when our training model requires only a subset of the TFRecord elements. For example, suppose that we have a multi-headed model that performs different types of pixel level segmentation on an input image. Each record in the data set contains multiple ground truth images corresponding to the multiple heads of the model. Now we decide to train a version of the model with just a single head. The trivial solution would be to feed in the complete TFRecord and simply ignore the unnecessary fields. However, the need to pull and parse extraneous data can severely impact the throughput of our input pipeline and thus the speed of our training. To reduce the risk of bottlenecks in the pipeline, we would need to create an additional copy of the data containing only the content relevant for the specific training at hand. This further exacerbates the data duplication issue we discussed above.
Record filtering: Sometimes we are interested in traversing records that have specific entry values. The trivial way to do this is to iterate over all of the data and simply drop any record that does not match our filter (e.g. using tf.filter). As before, this can introduce considerable overhead to the data input pipeline and result in significant starvation of the training loop. In previous posts (here and here) we demonstrated a need to filter input images based on whether they contained pink cars. The solution we proposed there was to store different data classes in separate files. Whenever the need to filter arose, we could traverse only the files associated with the filter. However, this solution requires us to be able to anticipate the filters we will need to apply which is not always possible. If we encounter a need to apply a filter that we did not anticipate we could either recreate the data with an additional partition or fall back to the trivial method of dropping non-matching samples. Neither option is optimal.
Data transformation: A typical input data pipeline might include multiple operations on the input data including data warping, augmentations, batching, and more. TensorFlow offers a set of built-in data processing operations that can be added to the input data pipeline computation graph via the tf.data.Dataset.map function. However, you might find that the data processing you require cannot be implemented efficiently using TensorFlow. One option would be to apply a block of native python code using tf.py_function, but this may limit your data throughput performance due to the Python Global Interpreter Lock (GIL). Many use cases call for using a format that enables applying operations in native python before being entered into the TensorFlow computation graph.
Based on the discussion above let us compile a list of some of the properties we are looking for in our file format. Given the added cost and complexities of storing multiple copies of our data, we will restrict ourselves from now on to a single copy of the data, i.e. we must choose a format that satisfies the needs of all data consumers. This list is not intended to be all inclusive. Additional requirements based on the particular use-case should be considered.
Distributed Storage: The file format and supporting libraries must support the option of storing large datasets in distributed storage settings such as Amazon S3, HDFS, and more.
Software Ecosystem: We require a strong ecosystem of libraries providing a range of tools for analyzing and manipulating the data.
Columnar Support: The file format and supporting libraries must support efficient extraction of a subset of features per data sample.
Row Filtering: The file format and supporting libraries should support efficient filtering on the values of sample features.
TensorFlow Integration: The format must facilitate efficient data streaming from storage into TensorFlow training sessions.
One requirement for training that we have left out of our discussion is data shuffling. It is common practice to shuffle the training data before each traversal (epoch). Were we able to randomly access any sample in the dataset, data shuffling would be easy. However, such random access to individual samples comes at the expense of performance, certainly in a distributed storage setting. Instead, one must resort to other mechanisms for shuffling including combinations of: shuffling samples during data creation, shuffling on the list of files that comprise the full dataset, interleaving between multiple datasets, and using an appropriately large shuffle buffer during training (e.g. see tf.data.Dataset.shuffle). We have chosen not to state shuffling as a requirement of the file format as, in our view, in the training scenario we have presented this challenge will exist with any choice of file format.
In the next section we will evaluate the compatibility of several file formats for training. In the following section we will discuss options for streaming them into TensorFlow.
A full survey of data formats for deep learning is beyond the scope of this post. We will limit our discussion to just a few strong candidates. More extensive surveys abound. Here is one example:
towardsdatascience.com
You can also check out a recent post of mine that covers the relatively new webdataset format, an intriguing solution specially geared towards developers who wish to maintain their data in the most raw form possible.
We will measure each format according to the metrics we chose above as in the below diagram in which we summarize our evaluation of the TFRecord format. We use green to indicate full support, yellow to indicate partial support, and red to indicate no support.
One of the more compelling file format options is Apache Parquet. Apache Parquet has an extensive software ecosystem with multiple frameworks and tools supporting a wide variety of data processing operations. It is especially popular among data analysts. The possibility of basing our training on the same format is quite attractive.
One of the main attributes of the Parquet format to which it owes much of its success is the fact that is a columnar storage format. Contrary to other formats such as CSV or TFRecord, in which each row of data is stored sequentially, in a columnar data format columns of data are stored together. Each file in the dataset is divided into row blocks each of which contains multiple samples. Within a row block the sample data is stored according to columns, i.e. the values of the first field of all of the samples appear first and are followed by the values of the second field of all of the samples, and so on. The columnar nature of the format facilitates efficient data analytics since queries can be made on subsets of columns without needing to load entire data records. Furthermore, grouping together columns can lead to more efficient compression of data and thus to reduced storage costs. Check out this post for more on the advantages of columnar data storage.
The following diagram summarizes our evaluation of the Parquet format:
Being a columnar format, Parquet enables efficient extraction of subsets of data columns. In addition, we can also take advantage of the columnar nature of the format to facilitate row filtering by: 1. first extracting the column on which we are filtering and then 2. extracting the rest of columns only for rows that match the filter. However, being that efficient data extraction might rely on extracting the columnar data of full blocks of rows, it is unclear that this method would perform well. Therefore we have marked this capability in yellow. A more performant filtering approach might require separating samples into different Parquet files according to class during the dataset creation, as described above. In pyspark this can be done using the partitionBy functionality.
Parquet Dataset Creation: In the code block below we demonstrate creation of a Parquet dataset from the popular Cifar10 dataset using the pyspark library. Parquet creation is supported by additional libraries as well including pandas and pyarrow.
from tensorflow.keras import datasetsfrom pyspark.sql import SparkSession, Rowfrom pyspark.sql.types import StructType, \        StructField, IntegerType, BinaryTypedef cifar_to_parquet():    schema = StructType(                [StructField("image", BinaryType(), True),                 StructField("label", IntegerType(), True)])    (data, labels), _ = datasets.cifar10.load_data()    labels = labels.flatten().tolist()    num_procs = 4 # set the number of parallel processes    spark = SparkSession.builder\        .master('local[{num_procs}]'.format(num_procs=num_procs))\        .getOrCreate()    sc = spark.sparkContext    num_samples = len(labels)    output_url = 'file:///tmp/parquet'    def row_generator(i):        return {            'image': bytearray(data[i].tobytes()),            'label': labels[i],        }    # optionally configure the size of row blocks    # blockSize = 1024 * 1024 * 16 # 16 MB    # sc._jsc.hadoopConfiguration()\    #     .setInt("parquet.block.size", blockSize)    rows_rdd = sc.parallelize(range(num_samples))\        .map(row_generator)\        .map(lambda x: Row(**x))    spark.createDataFrame(rows_rdd, schema)\        .write.mode('overwrite')\        .parquet(output_url)
While the number of files that comprise the dataset, the sizes of each file, and the size of each row block can have a meaningful impact on the performance of data loading, controlling these parameters can sometimes be a little tricky. In pyspark the row block size can be set using sc._jsc.hadoopConfiguration as in the comment in the code block above (the default value is 128 MB) and the sizes and number of files by the number of parallel processes and by the coalesce function.
Through appropriate configuration, the script can be modified to write directly to Amazon S3 (see e.g. here and here).
The petastorm library was created with the specific goal of unifying the dataset used by all data consumers on the development pipeline (see here). Although petastorm abstracts the underlying storage format, the default format used is Apache Parquet. More accurately, petastorm extends Apache Parquet by providing additional schema information in the Unischema structure, supporting multi-dimensional data, and supporting data compressing codecs. Henceforth we will assume the use of the extended Parquet as the underlying format and abuse the use of the term petastorm by using it to refer both to the library itself as well as to the file format created by the library. To distinguish between the two we will use capitalization when referring to the Petastorm format. The use of Parquet as the underlying format means that Petastorm enjoys all of the benefits associated with the use of the columnar format that we discussed above. However, use of the Parquet extensions introduced by petastorm does come with a caveat: The dataset needs to be created by the petastorm library and all data consumers on the development pipeline will require the petastorm library in order to properly read and parse the data. This requirement is somewhat limiting as one of the most appealing attributes of the Parquet format was its wide software ecosystem for accessing and manipulating data. It is for this reason that we have marked Software Ecosystem in yellow in the evaluation diagram below:
Petastorm includes a number of additional compelling features including support for row group indexing and n-grams. You can learn more about petastorm here and here.
Petastorm Dataset Creation: Dataset creation in petastorm is quite similar to Parquet creation in pyspark. The main differences are in the use of the Unischema and in the materialize_dataset context that wraps the dataset creation.
from petastorm.codecs import CompressedImageCodec, \        NdarrayCodec, ScalarCodecfrom petastorm.etl.dataset_metadata import materialize_datasetfrom petastorm.unischema import Unischema,\        UnischemaField, dict_to_spark_rowfrom pyspark.sql import SparkSessionfrom pyspark.sql.types import IntegerTypedef cifar_to_peta():    MySchema = Unischema('MySchema', [        UnischemaField('image', np.uint8,                       (32,32,3), NdarrayCodec(), False),        UnischemaField('label', np.uint8,                       (), ScalarCodec(IntegerType()), False),    ])    (data, labels), _ = datasets.cifar10.load_data()    labels = labels.flatten().tolist()    num_procs = 4 # set the number of parallel processes    spark = SparkSession.builder.\        master('local[{num_procs}]'.format(num_procs=num_procs))\        .getOrCreate()    sc = spark.sparkContext    num_samples = 100#len(labels)    output_url = 'file:///tmp/petastorm'    rowgroup_size_mb = 128    def row_generator(i):        return {            'image': data[i],            'label': np.uint8(labels[i]),        }    # Wrap dataset materialization portion.    # Will take care of setting up spark environment variables as    # well as save petastorm specific metadata    with materialize_dataset(spark, output_url,                             MySchema, rowgroup_size_mb):        rows_rdd = sc.parallelize(range(num_samples)) \            .map(row_generator) \            .map(lambda x: dict_to_spark_row(MySchema, x))        spark.createDataFrame(rows_rdd,                               MySchema.as_spark_schema()) \            .write \            .mode('overwrite') \            .parquet(output_url)
The row block size is determined in petastorm by the rowgroup_size_mb parameter. In this example we did not take advantage of the codec support included in petastorm. When using large data elements using the codec support could result in significant data compression and savings in storage costs.
The Feather file format is another columnar format which we consider due to its inclusion in TensorFlow I/O (more on this below). While there are many similarities between Feather and Parquet there are also a number of subtle differences resulting from their distinct underlying implementations. A good comparison between the two, including scenarios in which Feather might be the better option, can be found here. It is important to distinguish between version 1 and version 2 of the Feather format. Version 2 supports more data types as well as different types of data compression. Whenever you encounter a review of Feather, bear in mind that it may be based on version 1 of the format.
The primary difference between Feather and Parquet is in the extent of the software ecosystem. Although it is supported by libraries such as pyarrow and pandas, at the time of this writing the Feather format is far less popular than Parquet and the number of supporting frameworks is much more limited. We summarize our Feather file format evaluation in the diagram below:
Being a columnar format the same considerations as above exist for how we chose to rate the columnar support and row filtering properties.
Feather Dataset Creation: In the code block below we demonstrate creation of a Feather file using pyarrow. Feather creation support is built into the pandas library as well. Contrary to the previous dataset creations in which process parallelization was an integral part of the data creation code, here we demonstrate creation of a single file. A full solution would require spawning multiple processes responsible for creating disjoint subsets of the files comprising the full dataset.
from tensorflow.keras import datasetsimport pyarrow as pafrom pyarrow.feather import write_featherdef cifar_to_feather():    (data, labels), _ = datasets.cifar10.load_data()    data = [data[i].flatten() for i in range(data.shape[0])]    labels = labels.flatten()    table = pa.Table.from_arrays([data,labels], ['data','labels'])    write_feather(table, '/tmp/df.feather', chunksize=10000)write_feather(table, '/tmp/df.feather', chunksize=10000)
In the example above the size of the row block is determined by the chunksize parameter. Note that, contrary to the previous format creations, we determine the size by the number of records per block rather than the amount of memory per block.
We now turn our attention to the final requirement we listed above, compatibility of the file format with TensorFlow training. No matter how appealing your file format of choice may otherwise be, if you are not able to integrate its use into your training session, or if the speed of the input stream does not meet your needs, then you are back to square one. In this section we will explore some of the tools at our disposal for training in TensorFlow with the file formats we encountered above. In the following section we will measure the speed of the input stream in a number of experiments. Our discussion is based on versions 2.4.1 of TensorFlow, 0.17.1 of TensorFlow I/O, and 0.11.0 of petastorm.
In a typical TensorFlow application we define a tf.data.Dataset, which represents a sequence of data samples, and feed it into the training loop. Setting up a tf.data.Dataset includes defining the source of the data and applying transformations on the data. The source can be data stored in memory or in files. See here for more on how to create a dataset in TensorFlow. Since we have chosen the TFRecord format as our point of reference, let’s start by reviewing how we feed TFRecord files into a TensorFlow training session.
TFRecordDataset: A TFRecordDataset takes a list of TFRecord files as input and produces a sequence of serialized TFRecord data samples. It is typically followed by a tf.data.Dataset.map call in which each individual sample is parsed. The code block below demonstrates the creation of a TFRecordDataset from the Cifar10 data stored in TFRecord format:
import tensorflow as tfdef get_dataset():    autotune = tf.data.experimental.AUTOTUNE    def parse(example_proto):        feature_description = {            'image': tf.io.FixedLenFeature([], tf.string),            'label': tf.io.FixedLenFeature([], tf.int64)}        features = tf.io.parse_single_example(example_proto,                                                   feature_description)        image = tf.io.decode_raw(features['image'], tf.uint8)        image = tf.reshape(image, [32, 32, 3])        return image, label    records = tf.data.Dataset.list_files(<path_to_files>+'/*')    ds = tf.data.TFRecordDataset(records,                      num_parallel_reads=autotune)    ds = ds.map(parse, num_parallel_calls=autotune)    return ds
Our experience with TFRecordDataset has been overall positive; the underlying mechanics for pulling and parsing files seems to be pretty solid and we rarely find ourselves bottlenecked on this part of the pipeline. We have also found the API to be robust, performing well across a wide variety of training environments.
TensorFlow has several other tf.data.Dataset classes for loading data directly from files including FixedLengthRecordDataset and TextLineDataset. Building a TensorFlow dataset from a file format that doesn’t match any of the existing classes requires a bit more creativity. Here we will mention three options in ascending order of complexity.
Create the Dataset from memory source: The first option is to download and parse the files in Python (outside of TensorFlow) and load the data samples into TensorFlow from memory source. One way to do this is using tf.data.Dataset.from_generator as in the pseudo-code block below.
import tensorflow as tfdef get_custom_ds(file_names):    def my_generator():        for f in file_names:            # download f            samples = ... # parse samples from f            for sample in samples:                yield sample    return tf.data.Dataset.from_generator(                  my_generator,                  output_types=[tf.uint8,tf.uint8],                  output_shapes=[[32,32,3],[]])
Parse files in TensorFlow: A second option is to rely on TensorFlow’s tf.io module for downloading and parsing the files. In contrast with the previous solution, here the file management is part of the TensorFlow execution graph. Here is one way to do this using the tf.data.Dataset.list_files and tf.data.Dataset.interleave APIs:
import tensorflow as tfdef get_custom_ds():    autotune = tf.data.experimental.AUTOTUNE    filenames = tf.data.Dataset.list_files(<path_to_files>+'/*',                           shuffle=True)    def make_ds(path):        bytestring = tf.io.read_file(path)        samples = ... # parse bytestring using tf functions        return tf.data.Dataset.from_tensor_slices(samples)    ds = filenames.interleave(make_ds,                          num_parallel_calls=autotune,                         deterministic=False)    return ds
Create a custom Dataset class: The last option we mention is to create a new tf.data.Dataset class specially designed to handle your data format. This option requires the most technical prowess. It also offers the highest potential reward as measured by the speed of the data input stream. One way to implement this is to modify the TensorFlow C++ code and rebuild TensorFlow from source. For example, one could clone the TFRecordDataset implementation and overwrite only the portions of the code specifically relevant to parsing the format. In this way one would hope to be able to enjoy the same performance benefits as the TFRecordDataset. The disadvantage to this approach is that it will require maintaining a specialized version of TensorFlow. In particular, every time you upgrade to a new version of TensorFlow you will need to rebuild your custom solution. Note that custom Dataset class creation can also be implemented within TensorFlow I/O, rather than TensorFlow, as describe in this post.
While any of the above solutions can be appropriately tuned to maximize performance, it is not always that easy. To make matters worse, you may find that the ideal configurations (e.g. number of underling system processes) may vary greatly based on the training environment. In this sense using a dedicated Dataset class such as TFRecordDataset has a meaningful advantage over the customized solutions we have described. The next two solutions we will see will use Dataset classes specifically designed for our file format of choice.
TensorFlow I/O (tfio) is an extension package for TensorFlow that adds support for a number of file systems and file formats that are not included in TensorFlow. In particular, tfio defines the tfio.arrow.ArrowFeatherDataset class for creating datasets based on the Feather format and the tfio.v0.IODataset.from_parquet function for creating datasets based on the Parquet format.
TensorFlow I/O Feather Dataset: The tfio.arrow.ArrowFeatherDataset class is just one of a collection of APIs designed to support the Apache Arrow format. For a full overview of the tfio Apache Arrow offering, be sure to check out this blog. In the code block below we demonstrate the use of a tfio.arrow.ArrowFeatherDataset based on the Cifar10 data stored in Feather format that we created above.
import tensorflow as tfimport tensorflow_io.arrow as arrow_iodef get_dataset():    filenames = <list of feather files>    ds = arrow_io.ArrowFeatherDataset(filenames,                           columns=(0, 1),                          output_types=(tf.uint8, tf.uint8),                          output_shapes=([32*32*3,], []),                          batch_mode='auto')    ds = ds.unbatch()    return ds
By setting the batch_mode argument to ‘auto’ we are choosing to have the dataset return Parquet row blocks. Therefore the first call we apply is to unbatch the records so as to return individual samples. This strategy should result in better performance than reading samples individually.
We have found the throughput performance to improve if we combine the use of tfio.arrow.ArrowFeatherDataset with tf.data.Dataset.interleave:
import tensorflow as tfimport tensorflow_io as tfiodef get_dataset():    autotune = tf.data.experimental.AUTOTUNE    filenames = tf.data.Dataset.list_files(<path_to_files>+'/*',                           shuffle=True)    def make_ds(file):        ds = arrow_io.ArrowFeatherDataset(                     [file], [0,1],                     output_types=(tf.uint8, tf.uint8),                     output_shapes=([32*32*3,], []),                     batch_mode='auto')        return ds    ds = filenames.interleave(make_ds,                                     num_parallel_calls=autotune,                              deterministic=False)    ds = ds.unbatch()        return ds
TensorFlow I/O Parquet Dataset: In contrast with the Feather Dataset class, the from_parquet function receives a single Parquet file. However, we can overcome this limitation by using the tf.data.Dataset.interleave as shown below on the Cifar10 dataset stored in Parquet format:
import tensorflow as tfimport tensorflow_io as tfiodef get_dataset():    autotune = tf.data.experimental.AUTOTUNE    filenames = tf.data.Dataset.list_files(<path_to_files>+'/*',                           shuffle=True)    def parquet_ds(file):        ds = tfio.IODataset.from_parquet(file,                {'image': tf.string,                 'label': tf.int32})        return ds    ds = filenames.interleave(parquet_ds,                                     num_parallel_calls=autotune,                              deterministic=False)    def parse(example):        image = tf.io.decode_raw(example['image'], tf.uint8)        image = tf.reshape(image, [32, 32, 3])        label = example['label']        return image, label    ds = ds.map(parse,num_parallel_calls=autotune)        return ds
The petastorm library TensorFlow API defines the make_petastorm_dataset function for creating a TensorFlow tf.data.Dataset from a petastorm reader (petastorm.reader.Reader). The source of this dataset can be in either Petastorm format or raw Parquet format. To read from a dataset in Petastorm format we create the reader using the make_reader API. To read from a dataset in Parquet format we create the reader using the make_batch_reader API. There are a few delicate differences between the two readers described in the table here. Note that a TensorFlow tf.data.Dataset created from Petastorm format returns sequences of single data samples whereas a TensorFlow tf.data.Dataset created from raw Parquet format returns batches of data samples the size of which is determined by the Parquet row group size.
In the code block below we demonstrate the use of the make_petastorm_dataset API for creating a TensorFlow tf.data.Dataset from the Cifar10 data stored in Petastorm format.
from petastorm import make_readerfrom petastorm.tf_utils import make_petastorm_datasetdef get_dataset():    with make_reader('<path to data>') as reader:        ds = make_petastorm_dataset(reader)    return ds
In the code block below we demonstrate the use of the make_petastorm_dataset API for creating a TensorFlow tf.data.Dataset from the Cifar10 data stored in Parquet format.
from petastorm import make_batch_readerfrom petastorm.tf_utils import make_petastorm_datasetdef get_dataset():    autotune = tf.data.experimental.AUTOTUNE    with make_batch_reader('<path to data>') as reader:        ds = make_petastorm_dataset(reader)    ds = ds.unbatch()    def parse(example):        image, label = example        image = tf.io.decode_raw(image, tf.uint8)        image = tf.reshape(image, [32, 32, 3])        return image, label    ds = ds.map(parse,num_parallel_calls=autotune)    return ds
Note how we use the unbatch routine so as to return individual samples.
In this section we share the results of several experiments. All experiments were run on a c5.2xlarge Amazon EC2 instance type (which has 8 vCPUs) with TensorFlow version 2.4.1, TensorFlow I/O version 0.17.1, and petastorm version 0.11.0. The experiments were divided into two parts. First we experimented with different methods of feeding the Cifar10 data stored in the file formats we have discussed into a TensorFlow session. We created multiple copies of the data in order to artificially inflate the dataset. For these experiments we chose to set the training batch size to 1024.
In order to assess how the size of the sample records impacted the relative performance, we ran a second set of tests in which we added a random byte array, 2 MB in size, to each Cifar10 data sample. For these experiments we chose to set the training batch size to 16.
For all experiments the datasets were divided into underlying files of size 100–200 MB.
Since we were interested in measuring only the training data throughput, we chose to forgo building a training model and instead iterated directly on the TensorFlow dataset as in the following code block.
import timeds = get_dataset().batch(batch_size)round = 0start_time = time.time()for x in ds:    round = round + 1    if round % 100 == 0:        print("round {}: epoch time: {}".                format(round, time.time() - start_time))        start_time = time.time()    if round == 2000:        break
Note that in the case of petastorm the dataset traversal must be moved to within the petastorm reader context.
Our intention in sharing these results is to give you an idea of what your starting point might be and the amount of optimization effort that might be required. We strongly advise against drawing any conclusions from these results regarding your own use case for several reasons:
The performance of the solutions is likely to vary greatly based on the model, the dataset, and the training environment.For a given use case, the performance of each solution will vary based on the specifics of how both the format and the TensorFlow dataset are configured, including: the size of each file, the size of the row groups, the use of compression schemes, the number of worker processes, the size of each sample batch, and more.The experiments measure the maximum throughput of the solutions by iterating directly over the TensorFlow datasets. In practice, the throughput needs to be high enough to keep the training session busy. In other words, as long as we are not bottlenecked on IO, we might be perfectly content with a throughput that is lower than the maximum.By the time you read this post the versions of the libraries I have used might be outdated. Newer versions of the libraries might include changes that impact their performance.Last but not least is the likelihood of bugs in my experiments. When you find them please drop me a line.
The performance of the solutions is likely to vary greatly based on the model, the dataset, and the training environment.
For a given use case, the performance of each solution will vary based on the specifics of how both the format and the TensorFlow dataset are configured, including: the size of each file, the size of the row groups, the use of compression schemes, the number of worker processes, the size of each sample batch, and more.
The experiments measure the maximum throughput of the solutions by iterating directly over the TensorFlow datasets. In practice, the throughput needs to be high enough to keep the training session busy. In other words, as long as we are not bottlenecked on IO, we might be perfectly content with a throughput that is lower than the maximum.
By the time you read this post the versions of the libraries I have used might be outdated. Newer versions of the libraries might include changes that impact their performance.
Last but not least is the likelihood of bugs in my experiments. When you find them please drop me a line.
The results demonstrate the strength of the TFRecord format and TFRecordDataset. In particular, we can see their stability despite variations in the environment settings. At the same time, some of the other experiments resulted in training throughputs that do not pale in comparison. While we find the results to be encouraging, it is important to note the instability of the alternative solutions across the two cases we tested. Solutions that perform well (compared to the TFRecord format) in one scenario appear to break down in the other. It would seem that, in contrast with the use of the TFRecord format, adapting the other formats for use in varying training settings might require some extra effort. Given the advantages of the alternative formats that we have discussed in this post, this effort may be well worth its while.
An additional metric of interest is the CPU utilization associated with each of the options. In a recent blog post we discussed the possibility of bottlenecks in the training pipeline that are caused by reaching maximum utilization of one or more of the CPUs. In cases where this possibility is a concern, it would be wise to assess how different decisions related to the data storage affect the CPU utilization of loading and parsing the data. These include the choice of file format, data compression (and decompression), TensorFlow dataset creation method, and more. In the chart below we list the average CPU utilization for the three most performant solutions of the Cifar10 experiments. The CPU utilization was captured from the Amazon EC2 Instance Metrics in Amazon CloudWatch. Since the instance we used has 8 vCPUs, the maximum utilization is 800%.
Despite the clear favoritism of TensorFlow towards the TFRecord format, some of its properties leave it wanting. While the TFRecord format is great for training, it is unlikely to be suitable to any of the other consumers of the data on the development pipeline. In this post we have discussed several formats that better address the overall project needs and evaluated their compatibility to TensorFlow training. Although some effort may be required to support different training scenarios, our conclusion is that there are legitimate alternatives to the TFRecord format.