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 位於同一位置，並且將忽略不兼容的設備範圍。