tf.compat.v1.train.replica_device_setter

Возвращает функцию, которая используется при построении графа для реплик.

tf.compat.v1.train.replica_device_setter(
    ps_tasks=0,
    ps_device='/job:ps',
    worker_device='/job:worker',
    merge_devices=True,
    cluster=None,
    ps_ops=None,
    ps_strategy=None
)

Функции устройств используются в операторе with tf.device(device_function):, чтобы автоматически назначать устройства объектам Operation при их создании. Ограничения устройств добавляются из самого внутреннего контекста, работая наружу. Поведение слияния добавляет ограничения к полям, которые ещё не установлены более внутренним контекстом. В настоящее время полями являются (job, task, cpu/gpu).

Если cluster равно None, и ps_tasks равно 0, возвращаемая функция является функцией без действий. В противном случае, значение ps_tasks выводится из cluster.

По умолчанию, только операции переменных размещаются на задачах ps, и стратегия размещения выполняется по круговому циклу по всем задачам ps. Можно использовать пользовательскую функцию ps_strategy для более интеллектуального размещения, например, tf.contrib.training.GreedyLoadBalancingStrategy.

Например,

# To build a cluster with two ps jobs on hosts ps0 and ps1, and 3 worker
# jobs on hosts worker0, worker1 and worker2.
cluster_spec = {
    "ps": ["ps0:2222", "ps1:2222"],
    "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]}
with
tf.compat.v1.device(tf.compat.v1.train.replica_device_setter(cluster=cluster_spec)):
  # Build your graph
  v1 = tf.Variable(...)  # assigned to /job:ps/task:0
  v2 = tf.Variable(...)  # assigned to /job:ps/task:1
  v3 = tf.Variable(...)  # assigned to /job:ps/task:0
# Run compute