将 TensorFlow 梯度计算分成两个（或更多）部分

Question

是否可以部分使用 TensorFlow 的 tf.gradients() 函数，即 - 计算损失 w.r.t 一些张量的梯度，以及该张量 w.r.t 的权重，以及然后将它们相乘得到从损失到权重的原始梯度？

例如，设W,b为一些权重，设x为网络的输入，设y0表示标签。

假设正向图如

h=Wx+b
y=tanh(h)
loss=mse(y-y0)

我们可以计算tf.gradients(loss,W)然后应用（跳过一些细节）optimizer.apply_gradients()来更新W。

然后我尝试使用 var=tf.get_default_graph().get_tensor_by_name(...) 提取中间张量，然后计算两个梯度：g1=tf.gradients(loss,var) 和 g2=tf.gradients(var,W)。然后，根据链式法则，我希望 g1 和 g2 的维度能够计算出来，这样我就可以在某种意义上写出 g=g1*g2，然后返回 tf.gradients(loss,W)。

不幸的是，情况并非如此。尺寸不正确。每个渐变的尺寸都是 "w.r.t variable" 的尺寸，因此第一个渐变和第二个渐变之间不会存在对应关系。我缺少什么，我该怎么做？

谢谢。

Answer 1

来自 docs、tf.gradients（强调我的）

constructs symbolic derivatives of sum of ys w.r.t. x in xs.

如果 ys 中的任何张量是多维的，则在对标量的结果列表本身求和之前对其进行 reduce_summed，然后再进行微分。这就是为什么输出梯度与 xs.

具有相同大小的原因

这也解释了为什么在 tensorflow 中损失可以是多维的：它们在微分之前被隐含地求和。

Answer 2

tf.gradients 将对输入张量的梯度求和。为避免这种情况，您必须将张量拆分为标量并将 tf.gradients 应用于每个标量：

import tensorflow as tf

x = tf.ones([1, 10])

w = tf.get_variable("w", initializer=tf.constant(0.5, shape=[10, 5]))
out = tf.matmul(x, w)
out_target = tf.constant(0., shape=[5])

loss = tf.reduce_mean(tf.square(out - out_target))

grad = tf.gradients(loss, x)

part_grad_1 = tf.gradients(loss, out)
part_grad_2 = tf.concat([tf.gradients(i, x) for i in tf.split(out, 5, axis=1)], axis=1)

grad_by_parts = tf.matmul(part_grad_1, part_grad_2)

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)
    print(sess.run([grad]))
    print(sess.run([grad_by_parts]))

Answer 3

对于未来的读者：

Tensorflow 取得了一些进步，对于 tf2.7（甚至可能更早的版本）您可以使用 tf.GradientTape.jacobian 来避免对目标维度求和 .

https://www.tensorflow.org/guide/advanced_autodiff#jacobians

将 TensorFlow 梯度计算分成两个（或更多）部分

Breaking TensorFlow gradient calculation into two (or more) parts

automatic-differentiation

tensorflow