在 GPU 上选择性地注册一组操作的后向传递

Posted 2023-02-16

【中文标题】在 GPU 上选择性地注册一组操作的后向传递

【英文标题】：Selectively registering the backward pass of a set of ops on the GPU

【发布时间】：2015-12-24 22:57:23

【问题描述】：

在前向和后向（梯度）计算方面，我有一组在 CPU 上比在 GPU 上更快的操作。然而，它们只是整个模型的一小部分，其中大部分更适合在 GPU 上运行。目前，如果我在指定前向模型时只使用with tf.device(...)，并让 TF 决定将优化器放在哪里（例如tf.train.AdamOptimizer op），那么它将所有后向传递计算放在 GPU 上，这是次优的。有什么方法可以指定一个 op 及其梯度应该在 GPU 上注册吗？

【问题讨论】：

还有一个优化器最小化功能的 kwarg，colocate_gradients_with_ops

另一种方法是通过optimizer.compute_gradients显式指定要计算梯度的变量列表，然后将您的cpu计算梯度添加到列表中，并将计算梯度提供给optimizer.apply_gradients

【参考方案1】：

目前在（自动生成的）梯度计算中没有很好的方法来自定义操作的设备分配。但是，您可以做的一件事是使用with tf.device(): 注册“设备功能”（尽管this function 的文档适用并且更全面）。 “设备函数”是一个接受新构造的tf.Operation 并返回设备名称的函数，TensorFlow 将操作分配给该设备。这使您能够执行以下操作：

# These are almost certainly faster on GPU, but are just shown as an example.
OPS_ON_CPU = set(["AvgPool", "AvgPoolGrad"])

def _device_function(op):
  if op.type in OPS_ON_CPU:
    return "/cpu:0"
  else:
    # Other ops will be placed on GPU if available, otherwise CPU.
    return ""

with tf.device(_device_function):
  # Build model in here.
  # ...
  loss = ...
  train_op = tf.train.AdamOptimizer(0.01).minimize(loss) 

...这会将所有类型为 "AvgPool" 或 "AvgPoolGrad" 的操作放在 CPU 上。

【讨论】：

问题是我的操作并不慢，因为它们的 type，即它们只是由诸如 matmul 之类的相当普通的操作组成。但是它们的设置方式（大量的小型/精简计算）使它们在 CPU 上的速度更快（大约 3 倍）。

我几乎使用op.name 作为标准来完成这项工作，因为我的操作被包装在具有自己作用域的函数中。然而通过日志我注意到在梯度计算中有很多AddN ops仍然被放在GPU上（它们的名字像'gradients/AddN_x'），这仍然导致性能很差。跨度>

我终于得到了这个工作，但它需要更复杂的逻辑来测试一个操作的输入是否来自某个节点。 IE。而不是my_op in op.name 之类的条件，而是(my_op in op.name) or any(my_op in node.name for node in op.inputs)