tf.test.create_local_cluster

Создаёт и запускает локальные серверы и возвращает соответствующие Server объекты.

tf.test.create_local_cluster(
    num_workers,
    num_ps,
    protocol='grpc',
    worker_config=None,
    ps_config=None
)

"PS" означает "сервер параметров": задача, ответственная за хранение и обновление параметров модели. Другие задачи отправляют обновления этих параметров по мере работы над оптимизацией параметров. Такое разделение труда между задачами не является обязательным, но широко используется при распределённом обучении.

Подробнее читайте на https://www.tensorflow.org/guide/extend/architecture

Рисунок иллюстрирует взаимодействие этих компонентов. "/job:worker/task:0" и "/job:ps/task:0" — это две задачи со службами worker.

Пример:

workers, _ = tf.test.create_local_cluster(num_workers=2, num_ps=2)

worker_sessions = [tf.compat.v1.Session(w.target) for w in workers]

with tf.device("/job:ps/task:0"):
  ...
with tf.device("/job:ps/task:1"):
  ...
with tf.device("/job:worker/task:0"):
  ...
with tf.device("/job:worker/task:1"):
  ...

worker_sessions[0].run(...)