Python tf.data.experimental.load用法及代碼示例

加載以前保存的數據集。

用法

tf.data.experimental.load(
    path, element_spec=None, compression=None, reader_func=None
)

示例用法：

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)
for elem in new_dataset:
  print(elem)
tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)

請注意，要加載以前保存的數據集，您需要指定 element_spec - 已保存數據集元素的類型簽名，可以通過 tf.data.Dataset.element_spec 獲得。存在此要求，因此加載數據集的形狀推斷不需要執行 I/O。

如果使用了對已保存數據集進行分片的默認選項，則在加載時將保留已保存數據集的元素順序。

reader_func 參數可用於指定從各個分片加載元素的自定義順序。 reader_func 預計將采用單個參數——一個數據集的數據集，每個數據集包含一個分片的元素——並返回一個元素的數據集。例如，分片的順序可以在加載時打亂如下：

def custom_reader_func(datasets):
  datasets = datasets.shuffle(NUM_SHARDS)
  return datasets.interleave(lambda x:x, num_parallel_calls=AUTOTUNE)

dataset = tf.data.experimental.load(
    path="/path/to/data", ..., reader_func=custom_reader_func)