TensorFlow中如何获取多输出的独立梯度而非其求和梯度

好问题！确实，tf.gradients()（以及TF2中GradientTape.gradient()）默认会把所有输出求和后再计算对输入的梯度，但完全不用遍历每个输出就能一次性拿到所有独立梯度——核心思路是利用雅可比矩阵（Jacobian Matrix），它正好对应每个输出对每个输入的偏导数集合。

下面分TensorFlow版本给出具体实现：

TensorFlow 2.x（推荐） #

TF2的GradientTape提供了直接计算雅可比矩阵的方法，代码简洁高效，完全不需要循环：

import tensorflow as tf

# 定义10个输入变量
x = tf.Variable(tf.random.normal([10]))

# 记录前向传播
with tf.GradientTape() as tape:
    # 示例：生成5个输出（形状为[5]的张量）
    y = tf.matmul(tf.expand_dims(x, 0), tf.random.normal([10, 5]))[0]

# 计算雅可比矩阵：形状为[5, 10]，对应每个输出对每个输入的偏导
jacobian = tape.jacobian(y, x)
# 转置得到[10, 5]的数组，每一列对应一个输出对所有输入的梯度
gradients_matrix = tf.transpose(jacobian)

这里jacobian的形状是[输出维度, 输入维度]，转置后就完全符合你要的“10×5数组，每一列对应一个输出节点的梯度”的需求。

TensorFlow 1.x（图模式） #

如果还在维护TF1的老项目，可以利用tf.gradients()的grad_ys参数，传入单位矩阵作为权重，一次性生成所有独立梯度：

import tensorflow as tf

# 定义10个输入变量
x = tf.Variable(tf.random_normal([10]))
# 示例：生成5个输出（形状为[5]的张量）
y = tf.matmul(tf.expand_dims(x, 0), tf.random_normal([10, 5]))[0]

# 传入单位矩阵作为grad_ys，一次性计算所有输出的独立梯度
gradients_matrix = tf.gradients(y, x, grad_ys=tf.eye(5))[0]

tf.eye(5)生成一个5×5的单位矩阵，每个行向量对应给某个输出赋予权重1、其他输出权重0，这样tf.gradients就会返回每个输出对输入的梯度堆叠成的10×5矩阵，完美跳过遍历步骤。

原理补充 #

本质上，tf.gradients(y, x, grad_ys=w)计算的是sum(y_i * w_i)对x的梯度。当w是单位矩阵的第i行时，结果就是y_i对x的梯度；把所有单位行向量一起传入，就一次性得到了所有y_i对x的梯度集合，也就是雅可比矩阵的转置。

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