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