tf.distribute.coordinator.experimental_get_current_worker_index

Возвращает текущий индекс worker при вызове внутри замыкания worker.

tf.distribute.coordinator.experimental_get_current_worker_index()

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

Этот метод можно использовать внутри tf.function, выполняемой на worker. То есть внутри dataset_fn, выполняемого с помощью ClusterCoordinator.create_per_worker_dataset, или любой другой функции, запланированной с помощью ClusterCoordinator.schedule.

Пример (распределение данных по шардам в зависимости от worker):

strategy = tf.distribute.ParameterServerStrategy(
    cluster_resolver=...)
coordinator = (
    tf.distribute.coordinator.ClusterCoordinator(strategy))

def dataset_fn(context):
  dataset = tf.data.Dataset.range(10)
  worker_index = (
      tf.distribute.coordinator.experimental_get_current_worker_index()
  )
  dataset = dataset.shard(
      num_shards=num_workers,
      index=worker_index,
  )
  return dataset

@tf.function
def per_worker_dataset_fn():
  return strategy.distribute_datasets_from_function(dataset_fn)

per_worker_dataset = coordinator.create_per_worker_dataset(
    per_worker_dataset_fn)

Вызывает исключения
`RuntimeError`	если вызван вне `tf.function` или вне контекста удалённого выполнения замыкания (то есть на машине, не являющейся worker).

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/distribute/coordinator/experimental_get_current_worker_index