tensorflow - 多重损失的 tf 梯度

Question

假设我有一个多重损失定义为

losses = ... # a tensor with shape: (10,)

现在我想找到每个重量损失的梯度w

for i in range(len(10)):
    grad[i] = tf.gradients(losses[i], w)

现在，我可以在没有上述for循环的情况下直接执行此操作吗？

标签： tensorflow

Answer 1

您可以使用 tf.map_fn 在张量的第一维上映射任意函数。所以这样的事情应该可以解决问题：

def get_grads(x):
  return tf.gradients(x, w)

tf.map_fun(get_grads, losses)