TensorFlow设备规范

tf.DeviceSpec

DeviceSpec 类

定义在:tensorflow/python/framework/device.py.

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

在整个 TensorFlow 中使用 DeviceSpecs 来描述存储状态和计算发生的位置.使用 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):
  # 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)

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

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

DeviceSpec 由5个组件组成,每个组件都是可选指定的:

  • job:job 名称.
  • Replica:副本索引.
  • Task:任务索引.
  • Device type:设备类型字符串(例如 “CPU” 或 “GPU”).
  • Device index:设备索引.

属性

  • job
  • replica
  • task

方法

__init__

__init__(
    job=None,
    replica=None,
    task=None,
    device_type=None,
    device_index=None
)

创建一个新的 DeviceSpec 对象.

ARGS:

  • job:字符串,可选的 job 名称.
  • replica:int,可选的副本索引.
  • task:int,可选的任务索引.
  • device_type:可选的设备类型字符串(例如 “CPU” 或 “GPU”)
  • device_index:int,可选的设备索引.如果未指定,则 device 表示 “任意的” device_index.

from_string

from_string ( spec )

从字符串构造 DeviceSpec.

ARGS:

  • spec:一个字符串形式:/job:/replica:/task:/device:CPU:  或者 /job:/replica:/task:/device:GPU: 因为 cpu 和 gpu 是互斥的.所有条目都是可选的.

返回:

返回一个 DeviceSpec(设备规范).

merge_from

merge_from ( dev )

将 “dev” 的属性合并到此 DeviceSpec.

ARGS:

  • dev:一个 DeviceSpec.

parse_from_string

parse_from_string ( spec )

将 DeviceSpec 名称解析为其组件.

ARGS:

  • spec:一个字符串形式:/job:/replica:/task:/device:CPU:  或者 /job:/replica:/task:/device:GPU: 因为 cpu 和 gpu 是互斥的.所有条目都是可选的.

返回:

返回 DeviceSpec(设备规范).

注意:

  • ValueError:如果规范无效.

to_string

to_string ()

返回此 DeviceSpec 的字符串表示形式.

返回:

返回一个字符串的形式:/job:/replica:/task:/device::.


以上内容是否对您有帮助:
在线笔记
App下载
App下载

扫描二维码

下载编程狮App

公众号
微信公众号

编程狮公众号

意见反馈
返回顶部