tf.data.experimental.snapshot

API для сохранения результата входного набора данных.

tf.data.experimental.snapshot(
    path, compression='AUTO', reader_func=None, shard_func=None
)

API snapshot позволяет пользователям прозрачно сохранять выходные данные своей предварительной обработки на диск и материализовывать предварительно обработанные данные в другом запуске обучения.

Этот API позволяет объединить повторяющиеся шаги предварительной обработки и использовать уже обработанные данные, обменивая хранение на диске и пропускную способность сети на освобождение ценных ресурсов ЦП и ускорителей.

https://github.com/tensorflow/community/blob/master/rfcs/20200107-tf-data-snapshot.md содержит подробную документацию по проектированию этой функции.

Пользователи могут указать различные параметры для управления поведением snapshot, включая то, как считывать и записывать снимки, передавая пользовательские функции параметрам reader_func и shard_func.

shard_func — это пользовательская функция, которая сопоставляет элементы входных данных с фрагментами снимка.

Пользователи могут указать эту функцию, чтобы управлять тем, как файлы снимков записываются на диск. Ниже приведен пример того, как можно написать потенциальную функцию shard_func.

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

reader_func — это пользовательская функция, которая принимает один аргумент: (1) набор данных из наборов данных, каждый из которых представляет «раздел» элементов исходного набора данных. Мощность входного набора данных соответствует количеству фрагментов, указанных в 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.apply(tf.data.experimental.snapshot("/path/to/snapshot/dir",
    reader_func=user_reader_func))

По умолчанию snapshot параллелизует чтение по количеству ядер, доступных в системе, но не пытается перемешать данные.