tf.data.experimental.service.distribute

Преобразование, переносящее обработку набора данных в сервис tf.data.

tf.data.experimental.service.distribute(
    processing_mode,
    service,
    job_name=None,
    consumer_index=None,
    num_consumers=None,
    max_outstanding_requests=None,
    data_transfer_protocol=None,
    compression='AUTO',
    cross_trainer_cache=None,
    target_workers='AUTO'
) -> Callable[tf.data.Dataset, tf.data.Dataset]

При переборе набора данных, содержащего преобразование distribute, сервис tf.data создает «задание», которое формирует данные для перебора набора данных.

Сервис tf.data использует кластер рабочих процессов для подготовки данных к обучению модели. Аргумент processing_mode функции tf.data.experimental.service.distribute описывает, как задействовать несколько рабочих процессов для обработки входного набора данных. Сейчас доступны два режима обработки: «distributed_epoch» и «parallel_epochs».

Режим «distributed_epoch» означает, что набор данных будет разделен между всеми рабочими процессами сервиса tf.data. Диспетчер формирует «разделы» набора данных и отправляет их рабочим процессам для дальнейшей обработки. Например, если набор данных начинается со списка имен файлов, диспетчер переберет имена файлов и отправит их рабочим процессам tf.data, которые выполнят остальные преобразования набора данных для этих файлов. Режим «distributed_epoch» полезен, когда модели требуется увидеть каждый элемент набора данных ровно один раз или обработать данные в целом в последовательном порядке. Режим «distributed_epoch» работает только с наборами данных, источники которых можно разделить, например Dataset.from_tensor_slices, Dataset.list_files или Dataset.range.

Режим «parallel_epochs» означает, что каждый рабочий процесс сервиса tf.data независимо обработает весь входной набор данных. Поэтому важно перемешивать данные (например, имена файлов) недетерминированным образом, чтобы каждый рабочий процесс обрабатывал элементы набора данных в разном порядке. Режим «parallel_epochs» можно использовать для распределения наборов данных, которые нельзя разделить.

При наличии двух рабочих процессов режим «parallel_epochs» сформирует каждый элемент набора данных дважды:

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
# Start two workers
workers = [
    tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(
            dispatcher_address=dispatcher_address)) for _ in range(2)
]
dataset = tf.data.Dataset.range(10)
dataset = dataset.apply(tf.data.experimental.service.distribute(
    processing_mode="parallel_epochs", service=dispatcher.target))
print(sorted(list(dataset.as_numpy_iterator())))
[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9]

С другой стороны, режим «distributed_epoch» по-прежнему сформирует каждый элемент только один раз:

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
workers = [
    tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(
            dispatcher_address=dispatcher_address)) for _ in range(2)
]
dataset = tf.data.Dataset.range(10)
dataset = dataset.apply(tf.data.experimental.service.distribute(
    processing_mode="distributed_epoch", service=dispatcher.target))
print(sorted(list(dataset.as_numpy_iterator())))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

При использовании apply(tf.data.experimental.service.distribute(...)) набор данных до преобразования apply выполняется в сервисе tf.data, а операции после apply — в локальном процессе.

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
workers = [
    tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(
            dispatcher_address=dispatcher_address)) for _ in range(2)
]
dataset = tf.data.Dataset.range(5)
dataset = dataset.map(lambda x: x*x)
dataset = dataset.apply(
   tf.data.experimental.service.distribute("parallel_epochs",
                                           dispatcher.target))
dataset = dataset.map(lambda x: x+1)
print(sorted(list(dataset.as_numpy_iterator())))
[1, 1, 2, 2, 5, 5, 10, 10, 17, 17]

В приведенном выше примере операции с набором данных (до применения функции distribute к элементам) будут выполняться на рабочих процессах tf.data, а элементы будут передаваться по RPC. Остальные преобразования (после вызова distribute) будут выполняться локально. Диспетчер и рабочие процессы будут привязываться к незанятым портам (выбранным случайным образом), чтобы обмениваться данными. Чтобы привязать их к определенным портам, можно передать параметр port.

Аргумент job_name позволяет совместно использовать задания нескольким наборам данных. Вместо того чтобы каждый набор данных создавал собственное задание, все наборы данных с одинаковым job_name будут получать данные из одного задания. Для каждой итерации набора данных будет создано новое задание (каждый повтор Dataset.repeat считается новой итерацией). Предположим, что DispatchServer обслуживает запросы по адресу localhost:5000, а два рабочих процесса обучения (в конфигурации с одним или несколькими клиентами) перебирают приведенный ниже набор данных, при этом имеется один рабочий процесс tf.data:

range5_dataset = tf.data.Dataset.range(5)
dataset = range5_dataset.apply(tf.data.experimental.service.distribute(
    "parallel_epochs", "localhost:5000", job_name="my_job_name"))
for iteration in range(3):
  print(list(dataset))

Элементы каждого задания будут разделены между двумя процессами; процессы будут получать элементы в порядке поступления запросов. Например, процесс 1 может вывести:

[0, 2, 4]
[0, 1, 3]
[1]

а процесс 2 — вывести:

[1, 3]
[2, 4]
[0, 2, 3, 4]

Имена заданий нельзя повторно использовать для разных заданий обучения в течение срока работы сервиса tf.data. Как правило, предполагается, что сервис tf.data работает на протяжении одного задания обучения. Чтобы использовать сервис tf.data для нескольких заданий обучения, задавайте разные имена заданий во избежание конфликтов. Например, предположим, что задание обучения вызывает distribute с job_name="job" и считывает данные до конца входного потока. Если другое независимое задание подключится к тому же сервису tf.data и попытается читать из job_name="job", оно сразу получит признак конца входного потока, не получив никаких данных.

Координированное чтение данных

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

Например, сервис tf.data можно использовать для координации размеров примеров в кластере при синхронном обучении, чтобы на каждом шаге все реплики обучались на элементах схожего размера. Для этого определите набор данных, который формирует раунды по num_consumers последовательных пакетов схожего размера, а затем включите координированное чтение, задав consumer_index и num_consumers.

Примечание: Чтобы синхронизировать потребителей, при чтении данных по кругу набор данных должен иметь бесконечную кардинальность. Для этого можно добавить .repeat() в конец определения набора данных.

Keras и стратегии распределения

Набор данных, сформированный преобразованием distribute, можно передать в Keras Model.fit или в tf.distribute.Strategy.experimental_distribute_dataset стратегии распределения так же, как любой другой tf.data.Dataset. Рекомендуем задать job_name при вызове distribute, чтобы при наличии нескольких рабочих процессов они считывали данные из одного и того же задания. Обратите внимание, что автоматическое разделение на шарды, обычно выполняемое функцией experimental_distribute_dataset, будет отключено при задании job_name, поскольку совместное использование задания уже приводит к разделению данных между рабочими процессами. При совместном использовании задания данные будут динамически распределяться между рабочими процессами, чтобы они завершали чтение примерно одновременно. Это обеспечивает более эффективное использование рабочих процессов, чем автоматическое разделение на шарды, при котором каждый рабочий процесс обрабатывает независимый набор файлов, и у некоторых рабочих процессов данные могут закончиться раньше, чем у остальных.