多重损失的tf梯度

Question

假设我有多个损失定义为

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

现在，我想找到权重w上每个损失的梯度

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

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

Answer 1

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

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

tf.map_fun(get_grads, losses)