你能解释一下分布式 TensorFlow 教程示例吗？

Question

我对分布式计算的世界有些陌生。我正在阅读官方 tensorflow 教程中的 following，但我对教程的主要示例中发生的事情感到非常困惑。

特别是 ps 作业和工作者如何交互？ ps作业的作用究竟是什么？它们在代码中的对应部分非常有限，而且它们似乎没有做太多，那么它们的目的是什么？我想我不明白我们分布式系统的各个部分是如何协同工作的。

如果有人能根据不同的进程及其动作来解释你在最后执行 shell 命令时究竟发生了什么，那就太好了。

这里是主要代码供参考：

import argparse
import sys

import tensorflow as tf

FLAGS = None

def main(_):
  ps_hosts = FLAGS.ps_hosts.split(",")
  worker_hosts = FLAGS.worker_hosts.split(",")

  # Create a cluster from the parameter server and worker hosts.
  cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})

  # Create and start a server for the local task.
  server = tf.train.Server(cluster,
                           job_name=FLAGS.job_name,
                           task_index=FLAGS.task_index)

  if FLAGS.job_name == "ps":
    server.join()
  elif FLAGS.job_name == "worker":

    # Assigns ops to the local worker by default.
    with tf.device(tf.train.replica_device_setter(
        worker_device="/job:worker/task:%d" % FLAGS.task_index,
        cluster=cluster)):

      # Build model...
      loss = ...
      global_step = tf.contrib.framework.get_or_create_global_step()

      train_op = tf.train.AdagradOptimizer(0.01).minimize(
          loss, global_step=global_step)

    # The StopAtStepHook handles stopping after running given steps.
    hooks=[tf.train.StopAtStepHook(last_step=1000000)]

    # The MonitoredTrainingSession takes care of session initialization,
    # restoring from a checkpoint, saving to a checkpoint, and closing when done
    # or an error occurs.
    with tf.train.MonitoredTrainingSession(master=server.target,
                                           is_chief=(FLAGS.task_index == 0),
                                           checkpoint_dir="/tmp/train_logs",
                                           hooks=hooks) as mon_sess:
      while not mon_sess.should_stop():
        # Run a training step asynchronously.
        # See `tf.train.SyncReplicasOptimizer` for additional details on how to
        # perform *synchronous* training.
        # mon_sess.run handles AbortedError in case of preempted PS.
        mon_sess.run(train_op)

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.register("type", "bool", lambda v: v.lower() == "true")
  # Flags for defining the tf.train.ClusterSpec
  parser.add_argument(
      "--ps_hosts",
      type=str,
      default="",
      help="Comma-separated list of hostname:port pairs"
  )
  parser.add_argument(
      "--worker_hosts",
      type=str,
      default="",
      help="Comma-separated list of hostname:port pairs"
  )
  parser.add_argument(
      "--job_name",
      type=str,
      default="",
      help="One of 'ps', 'worker'"
  )
  # Flags for defining the tf.train.Server
  parser.add_argument(
      "--task_index",
      type=int,
      default=0,
      help="Index of task within the job"
  )
  FLAGS, unparsed = parser.parse_known_args()
  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

这里是shell命令：

  $ python trainer.py\
  --ps_hosts = ps0.example.com: 2222, ps1.example.com: 2222\
  --worker_hosts = worker0.example.com: 2222, worker1.example.com: 2222\
  --job_name = ps--task_index = 0# On ps1.example.com:
  $ python trainer.py\
  --ps_hosts = ps0.example.com: 2222, ps1.example.com: 2222\
  --worker_hosts = worker0.example.com: 2222, worker1.example.com: 2222\
  --job_name = ps--task_index = 1# On worker0.example.com:
  $ python trainer.py\
  --ps_hosts = ps0.example.com: 2222, ps1.example.com: 2222\
  --worker_hosts = worker0.example.com: 2222, worker1.example.com: 2222\
  --job_name = worker--task_index = 0# On worker1.example.com:
  $ python trainer.py\
  --ps_hosts = ps0.example.com: 2222, ps1.example.com: 2222\
  --worker_hosts = worker0.example.com: 2222, worker1.example.com: 2222\
  --job_name = worker--task_index = 1

Answer 1

这是情况的示意图。你有 4 个 tensorflow 进程。每个进程都运行 TensorFlow 工作线程，可以执行 TensorFlow 计算。此外，其中两个进程还运行一个客户端线程，该线程发出session.run 请求。

每个工作进程也是 TensorFlow 中的一个“设备”，目的是在设备上拆分图形执行。您可以通过在图形构建期间执行with tf.device("job:worker/task:0"): 之类的操作来告诉 TF 运行时在 worker1 设备上执行图形的某些部分。

tf.train.replica_device_setter 中发生了神奇的事情，它取代了手动的 with tf.device 注释，并具有跨设备自动分配变量的效果。更具体地说，当您有两个 PS 分片时，一半的变量将进入 ps1 设备，另一半进入 ps2 设备。同时，更新这些变量的图表部分将在每个工作设备上复制。

如果您将replica_device_setter 替换为手动设备规范，您的工作进程将大致如下所示

with tf.device('ps1'):
  var1 = tf.Variable(...)
with tf.device('ps2'):
  var2 = tf.Variable(...)
with tf.device('worker1'):
  update_op1 = var1.assign_add(grad1)
  update_op2 = var2.assign_add(grad2)

while True:
  sess.run([update_op1, update_op2])

通信会自动得到处理。当你在worker1客户端线程中执行sess.run(update_op1)时，它会在worker1上计算grad1，然后将结果发送给ps1任务，并触发ps1工作线程更新其var1的值

Answer 2

根据我的理解，ps作业包含了不同任务之间的所有共享数据，可以在不同的机器上运行（并且都共享同一个ps作业）。