Python tf.Graph.device用法及代码示例

用法

@tf_contextlib.contextmanager
device(
    device_name_or_function
)

返回指定要使用的默认设备的上下文管理器。

device_name_or_function 参数可以是设备名称字符串、设备函数或无：

• 如果它是设备名称字符串，则在此上下文中构造的所有操作都将分配给具有该名称的设备，除非被嵌套的 device() 上下文覆盖。
• 如果是函数，则将其视为从 Operation 对象到设备名称字符串的函数，并在每次创建新 Operation 时调用。操作将分配给具有返回名称的设备。
• 如果它是 None，则所有来自封闭上下文的 device() 调用都将被忽略。

有关设备名称字符串的有效语法的信息，请参阅 DeviceNameUtils 中的文档。

例如：

with g.device('/device:GPU:0'):
  # All operations constructed in this context will be placed
  # on GPU 0.
  with g.device(None):
    # All operations constructed in this context will have no
    # assigned device.

# Defines a function from `Operation` to device string.
def matmul_on_gpu(n):
  if n.type == "MatMul":
    return "/device:GPU:0"
  else:
    return "/cpu:0"

with g.device(matmul_on_gpu):
  # All operations of type "MatMul" constructed in this context
  # will be placed on GPU 0; all other operations will be placed
  # on CPU 0.

注意：设备范围可能会被操作包装器或其他库代码覆盖。例如，变量赋值操作 v.assign() 必须与 tf.Variable v 位于同一位置，并且将忽略不兼容的设备范围。