分布式 Tensorflow 如何处理这种 tf.Variable 创建?

How does distributed Tensorflow handle this kind of tf.Variable creation?

下面是分布式 Tensorflow 代码的两个版本,它们(尝试)实现一个存储在一个参数服务器上并由每个工作人员(异步)递增的全局计数器。

两个版本打印出来的好像是一样的东西,但是我不明白这是什么原因。版本之间的区别在于注释 # NEW 表示的两行。

当每个worker运行Version 1时,参数服务器是否自动为每个worker存储local_counter tf.Variables?

在版本 2 中,我尝试将每个 local_counter tf.Variable 明确地放在参数服务器上。

遵循版本 1 或版本 2 真的有区别吗?

PS:我确信这不是管理所有实例共享的 tf.Variable 的最佳方式,所以我很乐意接受任何改进建议。谢谢!

版本 1

# Standard distributed Tensorflow boilerplate
# ...

elif FLAGS.job_name == 'worker':

    TASK = FLAGS.task_index

    with tf.device('/job:ps/task:0/cpu:0'):
      with tf.variable_scope('global'):          
        global_counter = tf.Variable(0, name='global_counter',
                                        trainable=False)
        local_counter = tf.Variable(0, name='local_counter_{}'.format(TASK),
                                       trainable=False)
        init_op = tf.global_variables_initializer()

    with tf.device('/job:worker/task:{}'.format(TASK)): 
      with tf.variable_scope('local'):    
        local_inc_op = local_counter.assign_add(1) 
        global_inc_op = global_counter.assign_add(1)

    with tf.Session(server.target):
      sess.run(init_op)
      global_count = 0
      while global_count < 1000:
        sess.run([local_inc_op, global_inc_op])
        local_count, global_count = sess.run([local_counter, global_counter])
        print('Local {}, Global {}, worker-{}'.format(
               local_count, global_count, TASK))

版本 2

# Standard distributed Tensorflow boilerplate
# ...

elif FLAGS.job_name == 'worker':

    NUM_WORKERS = len(worker_hosts)
    TASK = FLAGS.task_index

    with tf.device('/job:ps/task:0/cpu:0'):
      with tf.variable_scope('global'):          
        global_counter = tf.Variable(0, name='global_counter',
                                        trainable=False)
        local_counters = [tf.Variable(0, name='local_counter_{}'.format(i),
                                          trainable=False)
                          for i in range(NUM_WORKERS)] # NEW
        init_op = tf.global_variables_initializer()

    with tf.device('/job:worker/task:{}'.format(TASK)): 
      with tf.variable_scope('local'):
        local_counter = local_counters[TASK] # NEW    
        local_inc_op = local_counter.assign_add(1) 
        global_inc_op = global_counter.assign_add(1)

    with tf.Session(server.target):
      sess.run(init_op)
      global_count = 0
      while global_count < 1000:
        sess.run([local_inc_op, global_inc_op])
        local_count, global_count = sess.run([local_counter, global_counter])
        print('Local {}, Global {}, worker-{}'.format(
               local_count, global_count, TASK))

我不确定我看到了多少实际差异。在任何一种情况下,本地计数器都是在参数服务 device 范围内创建的,因此它们将存在于参数服务器上。在版本 1 中,每个工作人员的图表仅包含其本地计数器,而每个工作人员的图表包含版本 2 中的所有本地计数器(但工作人员仍然只与他们自己的计数器交互,并且变量本身仍然存在于参数服务器上)。

所以要明确回答,是的,您可以将参数服务器图中不存在的变量存储在参数服务器上。基本上参数服务器(ResourceMgr)上有一个散列table,它可以存储任意变量names/values。

要自动将变量放在参数服务器上,tf.train.replica_device_setter 可以帮助减少样板文件。