Python tf.DeviceSpec用法及代码示例

表示 TensorFlow 设备的(可能是部分的)规范。

用法

tf.DeviceSpec(
    job=None, replica=None, task=None, device_type=None, device_index=None
)

DeviceSpec s 在整个 TensorFlow 中用于说明存储状态和发生计算的位置。使用DeviceSpec 允许您解析设备规范字符串以验证其有效性、合并它们或以编程方式组合它们。

例子：

# Place the operations on device "GPU:0" in the "ps" job.
device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
with tf.device(device_spec.to_string()):
  # Both my_var and squared_var will be placed on /job:ps/device:GPU:0.
  my_var = tf.Variable(..., name="my_variable")
  squared_var = tf.square(my_var)

在禁用 Eager Execution 的情况下(在 TensorFlow 1.x 中默认并在 TensorFlow 2.x 中调用 disable_eager_execution())，可以使用以下语法：

tf.compat.v1.disable_eager_execution()

# Same as previous
device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
# No need of .to_string() method.
with tf.device(device_spec):
  my_var = tf.Variable(..., name="my_variable")
  squared_var = tf.square(my_var)

如果部分指定了DeviceSpec，它将根据定义的范围与其他DeviceSpec 合并。 DeviceSpec 在内部作用域中定义的组件优先于在外部作用域中定义的组件。

gpu0_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
with tf.device(DeviceSpec(job="train").to_string()):
  with tf.device(gpu0_spec.to_string()):
    # Nodes created here will be assigned to /job:ps/device:GPU:0.
  with tf.device(DeviceSpec(device_type="GPU", device_index=1).to_string()):
    # Nodes created here will be assigned to /job:train/device:GPU:1.

DeviceSpec 由 5 个组件组成——每个组件都是可选的：

• 工作：工作名称。
• 副本：副本索引。
• 任务：任务索引。
• 设备类型：设备类型字符串(例如 "CPU" 或 "GPU")。
• 设备索引：设备索引。