Python tf.compat.v1.data.Dataset.snapshot用法及代碼示例

用法

snapshot(
    path, compression='AUTO', reader_func=None, shard_func=None, name=None
)

用於持久化輸入數據集輸出的 API。

快照 API 允許用戶透明地將其預處理管道的輸出保存到磁盤，並在不同的訓練運行中實現預處理數據。

此 API 可以整合重複的預處理步驟，並允許重新使用已處理的數據，以犧牲磁盤存儲和網絡帶寬來釋放更多寶貴的 CPU 資源和加速器計算時間。

https://github.com/tensorflow/community/blob/master/rfcs/20200107-tf-data-snapshot.md具有此函數的詳細設計文檔。

用戶可以指定各種選項來控製快照的行為，包括通過將用戶定義的函數傳遞給reader_func 和shard_func 參數來讀取和寫入快照的方式。

shard_func 是用戶指定的函數，用於將輸入元素映射到快照分片。

用戶可能希望指定此函數來控製快照文件應如何寫入磁盤。下麵是如何編寫潛在的shard_func 的示例。

dataset = ...
dataset = dataset.enumerate()
dataset = dataset.snapshot("/path/to/snapshot/dir",
    shard_func=lambda x, y:x % NUM_SHARDS, ...)
dataset = dataset.map(lambda x, y:y)

reader_func 是用戶指定的函數，它接受單個參數：(1) 數據集的數據集，每個數據集代表原始數據集元素的 "split"。輸入數據集的基數與 shard_func 中指定的分片數量相匹配(見上文)。該函數應返回原始數據集元素的數據集。

用戶可能希望指定此函數來控製應如何從磁盤讀取快照文件，包括混洗量和並行度。

這是用戶可以定義的標準閱讀器函數的示例。此函數啟用數據集洗牌和數據集的並行讀取：

def user_reader_func(datasets):
  # shuffle the datasets splits
  datasets = datasets.shuffle(NUM_CORES)
  # read datasets in parallel and interleave their elements
  return datasets.interleave(lambda x:x, num_parallel_calls=AUTOTUNE)

dataset = dataset.snapshot("/path/to/snapshot/dir",
    reader_func=user_reader_func)

默認情況下，快照按係統上可用的核心數並行讀取，但不會嘗試對數據進行混洗。