跨多个参数服务器的分片变量

正如我们很快会看到的那样，在分布式设置上训练神经网络时，常见模式是将模型参数存储在一组参数服务器上（即"ps"作业中的任务），而其他任务则集中在计算上（即 ，"worker"工作中的任务）。 对于具有数百万参数的大型模型，在多个参数服务器上分割这些参数非常有用，可以降低饱和单个参数服务器网卡的风险。 如果您要将每个变量手动固定到不同的参数服务器，那将非常繁琐。 幸运的是，TensorFlow 提供了replica_device_setter()函数，它以循环方式在所有"ps"任务中分配变量。 例如，以下代码将五个变量引入两个参数服务器：

with tf.device(tf.train.replica_device_setter(ps_tasks=2):    
    v1 = tf.Variable(1.0)  #  pinned to /job:ps/task:0    
    v2 = tf.Variable(2.0)  #  pinned to /job:ps/task:1    
    v3 = tf.Variable(3.0)  #  pinned to /job:ps/task:0    
    v4 = tf.Variable(4.0)  #  pinned to /job:ps/task:1    
    v5 = tf.Variable(5.0)  #  pinned to /job:ps/task:0

您不必传递ps_tasks的数量，您可以传递集群spec = cluster_spec，TensorFlow 将简单计算"ps"作业中的任务数。

如果您在块中创建其他操作，则不仅仅是变量，TensorFlow 会自动将它们连接到"/job:worker"，默认为第一个由"worker"作业中第一个任务管理的设备。 您可以通过设置worker_device参数将它们固定到其他设备，但更好的方法是使用嵌入式设备块。 内部设备块可以覆盖在外部块中定义的作业，任务或设备。 例如：

with tf.device(tf.train.replica_device_setter(ps_tasks=2)):    
    v1 = tf.Variable(1.0)  #  pinned to /job:ps/task:0 (+ defaults to /cpu:0)    
    v2 = tf.Variable(2.0)  #  pinned to /job:ps/task:1 (+ defaults to /cpu:0)    
    v3 = tf.Variable(3.0)  #  pinned to /job:ps/task:0 (+ defaults to /cpu:0)    
    [...]    
    s = v1 + v2            #  pinned to /job:worker (+ defaults to task:0/gpu:0)    
    with tf.device("/gpu:1"):        
        p1 = 2 * s         #  pinned to /job:worker/gpu:1 (+ defaults to /task:0)        
        with tf.device("/task:1"):            
            p2 = 3 * s     #  pinned to /job:worker/task:1/gpu:1

这个例子假设参数服务器是纯 CPU 的，这通常是这种情况，因为它们只需要存储和传送参数，而不是执行密集计算。

（未完成）