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

Question

我写了两个python文件create_save.py和load_use.py，如下所示。
create_save.py 运行良好并保存 tf 数据集。
但是 load_use.py 给出了如下所示的错误。
请如何修复 load_use.py 错误？
创建_保存.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

要加载以前保存的数据集，您需要指定 element_spec 参数 -- 保存数据集元素的类型签名，可以通过 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)。
在上面的例子中， element_spec load() 中的参数方法根据保存在代码中的数据的类型规范给出。
TF Data Load Documentation