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