固定任务的操作
通过指定作业名称,任务索引,设备类型和设备索引,可以使用设备块来锁定由任何任务管理的任何设备上的操作。 例如,以下代码将a
固定在"ps"
作业(即机器 A 上的 CPU)中第一个任务的 CPU,并将b
固定在"worker"
作业的第一个任务管理的第二个 GPU (这是 A 机上的 GPU# 1)。 最后,c
没有固定在任何设备上,所以主设备将它放在它自己的默认设备上(机器 B 的 GPU# 0 设备)。
with tf.device("/job:ps/task:0/cpu:0")
a = tf.constant(1.0)
with tf.device("/job:worker/task:0/gpu:1")
b = a + 2
c = a + b
如前所述,如果您省略设备类型和索引,则 TensorFlow 将默认为该任务的默认设备; 例如,将操作固定到"/job:ps/task:0"
会将其放置在"ps"
作业(机器 A 的 CPU)的第一个任务的默认设备上。 如果您还省略了任务索引(例如,"/job:ps"
),则 TensorFlow 默认为"/task:0"
。如果省略作业名称和任务索引,则 TensorFlow 默认为会话的主任务。