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

就向前和向后（渐变）计算而言，我有一组操作在CPU上比GPU更快。但是，它们只是整个模型的一小部分，其中大多数最好在GPU上运行。当前，如果我在指定正向模型时仅使用with tf.device(...)，并且让TF决定将优化器放置在何处（例如tf.train.AdamOptimizer op），那么它将所有反向传递计算放在GPU上，这是次优的。是否有某种方法可以指定应在GPU上注册op及其渐变？

最佳答案

当前，尚没有一种好的方法来为（自动生成的）梯度计算中的ops定制设备分配。但是，您可以做的一件事是使用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上。

关于tensorflow - 在GPU上选择性地注册一组操作的向后传递，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/34458603/