How can I wrap tf.io.parse_single_example with tf.py_function? - tensorflow

Question

First, I was wondering if I should wrap tf.io.parse_single_example with tf.py_function when reading TFRecord data from dataset.map,
N = config.get_num_listings_per_search()
features={
'qf': tf.io.FixedLenFeature([len(config.get_query_continuous_features())], tf.float32),
'qi': tf.io.FixedLenFeature([len(config.get_query_categorical_features())], tf.int64),
}
def _parse_function(example_proto):
parsed_features = tf.io.parse_single_example(example_proto, features)
return parsed_features['qf'], parsed_features['qi']
dataset = tf.data.TFRecordDataset(training_files)
dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
dataset = dataset.shuffle(buffer_size=1000000)
dataset = dataset.map(_parse_function, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.batch(config.get_batch_size())
becauseThe tf.data guide mentioned that
For performance reasons, we encourage you to use TensorFlow operations for preprocessing your data whenever possible. However, it is sometimes useful to call external Python libraries when parsing your input data. You can use the tf.py_function() operation in a Dataset.map() transformation.
I tried to wrap with
parsed_features = tf.py_function(tf.io.parse_single_example, (example_proto, features),
(tf.float32, tf.int64))
However, running the code gave me the following error:
TypeError: Tensors in list passed to 'input' of 'EagerPyFunc' Op have types [string, <NOT CONVERTIBLE TO TENSOR>] that are invalid.
It seems to me tf.py_function(tf.io.parse_single_example(example_proto, features)) is not supported because example_proto is of type tf.string ?
The primary reason I might want to do this is because the current input data pipeline is slow. Will I get some performance improvement if I wrap tf.io.parse_single_example with tf.py_function?
The above code is run in tensorflow-gpu==2.0
Thank you!

Answer 1

tf.py_function is meant to wrap external Python libraries like PIL or scipy, not TensorFlow operations like tf.io.parse_single_example. Adding tf.py_function here will probably make performance worse by forcing TensorFlow to call into Python instead of doing the parsing in C++.
The TFRecord guide gives an example of using tf.io.parse_single_example:
raw_image_dataset = tf.data.TFRecordDataset('images.tfrecords')
# Create a dictionary describing the features.
image_feature_description = {
'height': tf.io.FixedLenFeature([], tf.int64),
'width': tf.io.FixedLenFeature([], tf.int64),
'depth': tf.io.FixedLenFeature([], tf.int64),
'label': tf.io.FixedLenFeature([], tf.int64),
'image_raw': tf.io.FixedLenFeature([], tf.string),
}
def _parse_image_function(example_proto):
# Parse the input tf.Example proto using the dictionary above.
return tf.io.parse_single_example(example_proto, image_feature_description)
parsed_image_dataset = raw_image_dataset.map(_parse_image_function)