设备布置操作

TensorFlow 白皮书介绍了一种友好的动态布置器算法,该算法能够自动将操作分布到所有可用设备上,并考虑到以前运行图中所测量的计算时间,估算每次操作的输入和输出张量的大小, 每个设备可用的 RAM,传输数据进出设备时的通信延迟,来自用户的提示和约束等等。 不幸的是,这种复杂的算法是谷歌内部的,它并没有在 TensorFlow 的开源版本中发布。它被排除在外的原因似乎是,由用户指定的一小部分放置规则实际上比动态放置器放置的更有效。 然而,TensorFlow 团队正在努力改进它,并且最终可能会被开放。

在此之前,TensorFlow都是简单的放置,它(如其名称所示)非常基本。

简单放置

无论何时运行图形,如果 TensorFlow 需要求值尚未放置在设备上的节点,则它会使用简单放置器将其放置在未放置的所有其他节点上。 简单放置尊重以下规则:

  • 如果某个节点已经放置在图形的上一次运行中的某个设备上,则该节点将保留在该设备上。
  • 否则,如果用户将一个节点固定到设备上(下面介绍),则放置器将其放置在该设备上。
  • 否则,它默认为 GPU# 0,如果没有 GPU,则默认为 CPU。

正如您所看到的,将操作放在适当的设备上主要取决于您。 如果您不做任何事情,整个图表将被放置在默认设备上。 要将节点固定到设备上,您必须使用device()函数创建一个设备块。 例如,以下代码将变量a和常量b固定在 CPU 上,但乘法节点c不固定在任何设备上,因此将放置在默认设备上:

  1. with tf.device("/cpu:0"):
  2. a = tf.Variable(3.0)
  3. b = tf.constant(4.0)
  4. c = a * b

其中,"/cpu:0"设备合计多 CPU 系统上的所有 CPU。 目前没有办法在特定 CPU 上固定节点或仅使用所有 CPU 的子集。

记录放置位置

让我们检查一下简单的放置器是否遵守我们刚刚定义的布局约束条件。 为此,您可以将log_device_placement选项设置为True;这告诉放置器在放置节点时记录消息。例如:

  1. >>> config = tf.ConfigProto()
  2. >>> config.log_device_placement = True
  3. >>> sess = tf.Session(config=config)
  4. I [...] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GRID K520, pci bus id: 0000:00:03.0)
  5. [...]
  6. >>> x.initializer.run(session=sess)
  7. I [...] a: /job:localhost/replica:0/task:0/cpu:0
  8. I [...] a/read: /job:localhost/replica:0/task:0/cpu:0
  9. I [...] mul: /job:localhost/replica:0/task:0/gpu:0
  10. I [...] a/Assign: /job:localhost/replica:0/task:0/cpu:0
  11. I [...] b: /job:localhost/replica:0/task:0/cpu:0
  12. I [...] a/initial_value: /job:localhost/replica:0/task:0/cpu:0
  13. >>> sess.run(c)
  14. 12

Info中以大写字母I开头的行是日志消息。 当我们创建一个会话时,TensorFlow 会记录一条消息,告诉我们它已经找到了一个 GPU 卡(在这个例子中是 Grid K520 卡)。 然后,我们第一次运行图形(在这种情况下,当初始化变量a时),简单布局器运行,并将每个节点放置在分配给它的设备上。正如预期的那样,日志消息显示所有节点都放在"/cpu:0"上,除了乘法节点,它以默认设备"/gpu:0"结束(您可以先忽略前缀:/job:localhost/replica:0/task:0;我们将在一会儿讨论它)。 注意,我们第二次运行图(计算c)时,由于 TensorFlow 需要计算的所有节点c都已经放置,所以不使用布局器。

动态放置功能

创建设备块时,可以指定一个函数,而不是设备名称。TensorFlow 会调用这个函数来进行每个需要放置在设备块中的操作,并且该函数必须返回设备的名称来固定操作。 例如,以下代码将固定所有变量节点到"/cpu:0"(在本例中只是变量a)和所有其他节点到"/gpu:0"

  1. def variables_on_cpu(op):
  2. if op.type == "Variable":
  3. return "/cpu:0"
  4. else:
  5. return "/gpu:0"
  6. with tf.device(variables_on_cpu):
  7. a = tf.Variable(3.0)
  8. b = tf.constant(4.0)
  9. c = a * b

您可以轻松实现更复杂的算法,例如以循环方式用GPU锁定变量。

操作和内核

对于在设备上运行的 TensorFlow 操作,它需要具有该设备的实现;这被称为内核。 许多操作对于 CPU 和 GPU 都有内核,但并非全部都是。 例如,TensorFlow 没有用于整数变量的 GPU 内核,因此当 TensorFlow 尝试将变量i放置到 GPU# 0 时,以下代码将失败:

  1. >>> with tf.device("/gpu:0"):
  2. ... i = tf.Variable(3)
  3. [...]
  4. >>> sess.run(i.initializer)
  5. Traceback (most recent call last):
  6. [...]
  7. tensorflow.python.framework.errors.InvalidArgumentError: Cannot assign a device to node 'Variable': Could not satisfy explicit device specification

请注意,TensorFlow 推断变量必须是int32类型,因为初始化值是一个整数。 如果将初始化值更改为 3.0 而不是 3,或者如果在创建变量时显式设置dtype = tf.float32,则一切正常。

软放置

默认情况下,如果您尝试在操作没有内核的设备上固定操作,则当 TensorFlow 尝试将操作放置在设备上时,您会看到前面显示的异常。 如果您更喜欢 TensorFlow 回退到 CPU,则可以将allow_soft_placement配置选项设置为True

  1. with tf.device("/gpu:0"):
  2. i = tf.Variable(3)
  3. config = tf.ConfigProto()
  4. config.allow_soft_placement = True
  5. sess = tf.Session(config=config)
  6. sess.run(i.initializer) # the placer runs and falls back to /cpu:0

到目前为止,我们已经讨论了如何在不同设备上放置节点。 现在让我们看看 TensorFlow 如何并行运行这些节点。