如何使用实验性的保存和加载方法保存和使用Tensorflow数据集？

Question

我编写了两个python文件create_save.py和load_use.py，如下所示。 create_save.py运行良好，正在保存tf数据集

但是load_use.py给出的错误如下所示。 如何修复load_use.py错误

创建_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)

上述load_use.py程序出现以下错误：

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'

如何修复上述错误

Answer 1

要加载以前保存的数据集，需要指定保存的数据集元素的类型签名，该签名可通过tf.data.dataset.element_spec获得。此要求的存在使得加载的数据集的形状推断不需要执行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)

创建tf.data.Dataset时，它具有属性element_spec，这是加载保存的文件时应该使用的属性。（请参阅：Dataset doc）

在上面的示例中，load()方法中的element_spec参数是根据代码中保存的数据的类型规范给出的

TF Data Load Documentation