TensorFlow中未知尺寸张量的雅可比矩阵计算及图构建阶段梯度求解

在处理带None维度（比如动态batch size）的张量时，直接用tf.gradients只能得到求和后的梯度，没法获取每个元素单独的梯度（也就是雅可比矩阵）。不过我们可以通过tf.map_fn遍历张量的动态维度，逐个计算梯度，完美解决这个问题。

针对你给出的代码场景：

x = tf.placeholder("float", [None,3])
v = tf.Variable(tf.ones([3, 1]))
w = tf.matmul(x,v)

这里w的形状是[None, 1]，我们需要计算w中每个元素对v的梯度，最终得到的雅可比矩阵形状应该是[None, 3]（因为每个w_i是x的第i行与v的点积，对v的梯度就是x的第i行）。

具体实现步骤 #

1. 先把w压缩成一维张量（去掉最后一个维度的1），方便遍历：

   w_squeezed = tf.squeeze(w, axis=1)  # 形状变为[None]
   

2. 用tf.map_fn遍历w_squeezed的每个元素，分别计算对v的梯度：

   # 对每个w元素计算梯度，map_fn会自动适配动态的None维度
   jacobian = tf.map_fn(lambda elem: tf.gradients(elem, v)[0], w_squeezed, dtype=tf.float32)
   # 把结果中的多余维度去掉，得到[None, 3]的雅可比矩阵
   jacobian = tf.squeeze(jacobian, axis=2)
   

完整测试代码 #

import tensorflow as tf

# 构建计算图
x = tf.placeholder(tf.float32, [None, 3])
v = tf.Variable(tf.ones([3, 1]))
w = tf.matmul(x, v)
w_squeezed = tf.squeeze(w, axis=1)

# 计算雅可比矩阵
jacobian = tf.map_fn(lambda elem: tf.gradients(elem, v)[0], w_squeezed, dtype=tf.float32)
jacobian = tf.squeeze(jacobian, axis=2)

# 运行测试
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 输入一个batch size为2的x
    x_input = [[1, 2, 3], [4, 5, 6]]
    jac_val, w_val = sess.run([jacobian, w], feed_dict={x: x_input})
    
    print("w的计算结果:\n", w_val)
    print("雅可比矩阵（w每个元素对v的梯度）:\n", jac_val)

结果说明 #

运行后你会发现，雅可比矩阵的每一行正好是输入x的对应行——这完全符合数学推导：因为w_i = x_i1*v1 + x_i2*v2 + x_i3*v3，对v的梯度就是[x_i1, x_i2, x_i3]，和我们的实现结果一致。

这种方法的优势是在构建计算图阶段不需要知道None维度的具体数值，TensorFlow会在运行时根据输入的实际batch size自动处理，完美适配动态尺寸的场景。

内容的提问来源于stack exchange，提问作者Zhiwei