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 параллелизует чтение по количеству ядер, доступных в системе, но не будет пытаться перемешать данные.