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

Question

我有一组运算在 CPU 上比在 GPU 上更快，无论是前向还是后向(梯度)计算。然而，它们只是整个模型的一小部分，其中大部分在 GPU 上运行效果更好。目前，如果我只是在指定正向模型时使用 with tf.device(...)，并且让 TF 决定优化器的放置位置(例如 tf.train.AdamOptimizer op)，然后它将所有反向传递计算放在 GPU 上，这是次优的。是否有某种方法可以指定应在 GPU 上注册操作及其梯度？

Answer 1

目前，在(自动生成的)梯度计算中，没有很好的方法来为 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 上。