TensorFlow GradientTape目标为非标量时的输出形态与取值解析

当你用GradientTape.gradient()处理非标量的target时，TensorFlow默认会先对target的所有元素求和，再计算这个总和对自变量的梯度——这就是你看到grads结果的原因。

结合你的代码示例分析 #

先看代码里的数值计算：

• a是形状为(3,1)的变量，b是(2,3)的变量，矩阵乘法b @ a得到的c是(2,1)的张量：[[2.],[2.]]
• 对c的所有元素求和得到4.，然后计算这个总和对a的梯度，就得到了输出的[[5.],[7.],[9.]]

手动推导验证一下：
总和S = c[0][0] + c[1][0] = (1*(-1)+2*0+3*1) + (4*(-1)+5*0+6*1) = 2 + 2 = 4
对a的每个元素求偏导：

• a[0][0]的偏导：1（来自c[0][0]） + 4（来自c[1][0]）= 5
• a[1][0]的偏导：2 + 5 =7
• a[2][0]的偏导：3 +6=9
  完全匹配输出的梯度结果。

获取每个target元素的偏导数（雅可比矩阵） #

如果你想得到c中每个元素对a的偏导数，需要用GradientTape.jacobian()方法，而不是gradient()。示例代码如下：

import tensorflow as tf

a = tf.Variable([[-1.], [0.], [1.]])
b = tf.Variable([[1.,2.,3.],[4.,5.,6.]])
with tf.GradientTape() as g:
    c = b @ a
# 计算雅可比矩阵，原始形状为(2,1,3,1)
jacobian = g.jacobian(c, a)
# 整理为更直观的(2,3)形状
jacobian_reshaped = tf.reshape(jacobian, (2,3))
print("雅可比矩阵：")
print(jacobian_reshaped)

运行后输出：

雅可比矩阵：
[[1. 2. 3.]
 [4. 5. 6.]]

这个矩阵的每一行对应c的一个元素对a所有元素的偏导数，比如第一行是c[0][0]对a三个元素的偏导，第二行是c[1][0]对a三个元素的偏导。

核心结论 #

• gradient()方法：默认仅对标量target计算梯度，若target是非标量，会自动对其所有元素求和后计算梯度，输出形状与自变量一致
• jacobian()方法：计算每个target元素对自变量的偏导数，输出形状为target形状 + 自变量形状，可通过reshape调整为更直观的形态

内容的提问来源于stack exchange，提问作者Joe C.