定义服务器和设备之间的参数和操作
您可以使用我们在第 1 章中使用的tf.device()
函数,将参数放在ps
任务和worker
任务上图的计算节点上。
请注意,您还可以通过将设备字符串添加到任务字符串来将图节点放置在特定设备上,如下所示:/job:<job name>/task:<task index>/device:<device type>:<device index>
.
对于我们的演示示例,我们使用 TensorFlow 函数tf.train.replica_device_setter()
来放置变量和操作。
- 首先,我们将 worker 设备定义为当前 worker:
worker_device='/job:worker/task:{}'.format(FLAGS.task_index)
- 接下来,使用
replica_device_setter
定义设备功能,传递集群规范和当前工作设备。replica_device_setter
函数从集群规范中计算出参数服务器,如果有多个参数服务器,则默认情况下以循环方式在它们之间分配参数。参数放置策略可以更改为tf.contrib
包中的用户定义函数或预构建策略。
device_func = tf.train.replica_device_setter(
worker_device=worker_device,cluster=clusterSpec)
- 最后,我们在
tf.device(device_func)
块内创建图并训练它。对于同步更新和异步更新,图的创建和训练是不同的,因此我们将在两个单独的小节中介绍这些内容。