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")。
• 設備索引：設備索引。