TensorFlow操作的设备分配：能否在创建后指定运行设备？

1. 通过属性为TensorFlow操作分配设备 #

有两种常用的方式可以在创建操作时指定它运行的设备：

• 使用tf.device()上下文管理器：这是最通用的方法，把需要分配到特定设备的操作代码包裹在上下文里，该范围内的所有操作都会被分配到指定设备（前提是你的硬件支持该设备）。示例代码：

import tensorflow as tf

with tf.device('/gpu:1'):
    # 这个上下文里的所有操作都会被分配到/gpu:1
    input_tensor = tf.constant([[1.0, 2.0], [3.0, 4.0]])
    output_tensor = tf.matmul(input_tensor, input_tensor)

• 直接在操作构造时指定device参数：部分TensorFlow操作的构造函数支持device参数，可以直接给单个操作指定设备。示例：

# 直接将这个常量操作分配到/gpu:1
const_op = tf.constant(10.0, device='/gpu:1')

注意：不是所有操作都支持这个参数，大部分场景下还是推荐用上下文管理器的方式。

2. 操作创建后能否修改其设备分配？ #

很遗憾，不能直接修改已创建操作的设备。TensorFlow的Operation对象（也就是你说的op）在创建完成后，其device属性是只读的，你尝试用op.device = '/gpu:1'这种方式修改会报错。

这是因为TensorFlow的计算图是静态结构，操作一旦被添加到图中，它的设备分配就已经确定（要么是你手动指定的，要么是TensorFlow的自动设备 placer自动分配的），无法再变更。

替代方案 #

如果需要让逻辑相同的操作运行在不同设备上，只能在目标设备的上下文里重新创建对应的操作。你可以把操作的定义封装成函数，这样就能方便地在不同设备上下文里复用：

import tensorflow as tf

def create_matmul_op():
    input_tensor = tf.constant([[1.0, 2.0], [3.0, 4.0]])
    return tf.matmul(input_tensor, input_tensor)

# 在默认设备上创建操作
op_default = create_matmul_op()
# 在/gpu:1上重新创建相同逻辑的操作
with tf.device('/gpu:1'):
    op_gpu1 = create_matmul_op()

如果是在Eager Execution模式下，逻辑也是一样的：操作一旦执行（创建），设备就固定了，只能重新创建来切换设备。

内容的提问来源于stack exchange，提问作者Ahmad