How to save and use a Tensorflow dataset using the Experimental save and load mehods?

Question

I wrote two python files create_save.py and load_use.py as shown below. create_save.py is running good and saving tf dataset.

But load_use.py is giving errors shown below. How to fix load_use.py errors please?

create_save.py

import os
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.data.experimental import save as tf_save
ds_dir = os.path.join('./', "save_load_tfds_dir")
ds = tf.data.Dataset.range(12)
tf_save(ds, ds_dir)

load_use.py

import os
import numpy as np
import pandas as pd
import tensorflow as tf
ds_dir = os.path.join('./', "save_load_tfds_dir")
new_ds = tf.data.experimental.load(ds_dir)
for elem in new_ds:
  print(elem)

Above load_use.py program is giving following errors:

TypeError Traceback (most recent call last) in ----> 1 new_ds = tf.data.experimental.load(ds_dir)

TypeError: load() missing 1 required positional argument: 'element_spec'

How to fix above error?

Answer 1

To load a previously saved dataset, you need to specify element_spec argument -- a type signature of the elements of the saved dataset, which can be obtained via tf.data.Dataset.element_spec. This requirement exists so that shape inference of the loaded dataset does not need to perform I/O.

import tempfile
path = os.path.join(tempfile.gettempdir(), "saved_data")
# Save a dataset
dataset = tf.data.Dataset.range(2)
tf.data.experimental.save(dataset, path)
new_dataset = tf.data.experimental.load(path,
    tf.TensorSpec(shape=(), dtype=tf.int64))  # element_spec arg
for elem in new_dataset:
  print(elem)

When you are creating a tf.data.Dataset, it has the attribute element_spec which is what you should be using while loading your saved file. (Refer: Dataset doc).

In the above example, the element_spec argument in the load() method is given as per the type specification of the data being saved in the code.

TF Data Load Documentation