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" — это две задачи с рабочими службами.

Пример:

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(...)