TensorFlow中计算dy/dx梯度：需保留y+x维度，解决tf.gradients求和问题

针对你提出的两个梯度计算需求，我来拆解一下具体的实现方案：

1. 计算与y尺寸一致的dy/dx #

你提到tf.gradients返回的是sum(dy/dx)而非dy/dx本身，这是因为它默认会对y的所有元素求和后再计算对x的梯度，所以结果维度和x一致。如果想要得到和y尺寸完全相同的梯度结果，我们需要针对y的每个元素单独处理，或者对梯度的x维度做聚合：

TensorFlow 2.x 实现 #

用tf.GradientTape的jacobian方法可以直接拿到每个y元素对x的梯度，再根据需求聚合到y的维度：

import tensorflow as tf

# 示例：y = A@x，y.shape=(100,1)，x.shape=(50,1)
x = tf.Variable(tf.random.normal((50, 1)))
A = tf.random.normal((100, 50))

with tf.GradientTape() as tape:
    y = tf.matmul(A, x)

# 先计算完整雅可比矩阵（100,1,50,1），再对x的维度求和得到与y同尺寸的结果
jacobian = tape.jacobian(y, x)
dy_dx = tf.reduce_sum(jacobian, axis=[2, 3])
print(dy_dx.shape)  # 输出 (100, 1)，和y尺寸一致

如果x是标量，直接用jacobian就能得到和y同维度的梯度：

x = tf.Variable(2.0)
with tf.GradientTape() as tape:
    y = tf.stack([x**2, x**3])  # y.shape=(2,)

dy_dx = tape.jacobian(y, x)
print(dy_dx.shape)  # 输出 (2,)，匹配y的尺寸

TensorFlow 1.x 实现 #

在TF1.x中可以用tf.map_fn遍历y的每个元素，逐个计算梯度后再处理维度：

import tensorflow as tf

x = tf.Variable(tf.random.normal((50, 1)))
A = tf.random.normal((100, 50))
y = tf.matmul(A, x)  # (100,1)

# 定义单个y元素对x的梯度计算函数
def single_element_grad(y_element):
    return tf.gradients(y_element, x)[0]

# 遍历y的所有元素，得到每个元素对x的梯度，再求和压缩维度
per_element_grads = tf.map_fn(single_element_grad, tf.squeeze(y))
dy_dx = tf.expand_dims(tf.reduce_sum(per_element_grads, axis=1), axis=1)
print(dy_dx.shape)  # 输出 (100,1)

2. 计算包含x维度的完整梯度（y尺寸 + x尺寸） #

这个需求其实就是计算雅可比矩阵——y中每个元素对x中每个元素的偏导数，结果维度正好是y.shape + x.shape，完全匹配你给出的[100x1x50x1]示例：

TensorFlow 2.x 实现 #

TF2.x的GradientTape.jacobian方法直接支持这个需求，一步到位：

x = tf.Variable(tf.random.normal((50, 1)))
A = tf.random.normal((100, 50))

with tf.GradientTape() as tape:
    y = tf.matmul(A, x)

jacobian = tape.jacobian(y, x)
print(jacobian.shape)  # 输出 (100, 1, 50, 1)，完美符合要求

如果是批量场景，还可以用batch_jacobian来高效处理批量数据，性能会比循环好很多。

TensorFlow 1.x 实现 #

TF1.x没有原生的雅可比函数，需要手动循环堆叠梯度：

x = tf.Variable(tf.random.normal((50, 1)))
A = tf.random.normal((100, 50))
y = tf.matmul(A, x)  # (100,1)

grads_list = []
for i in range(y.shape[0]):
    # 计算第i个y元素对x的梯度
    grad = tf.gradients(y[i], x)[0]
    # 增加维度以便堆叠
    grads_list.append(tf.expand_dims(grad, axis=0))

# 堆叠后调整到目标维度
jacobian = tf.stack(grads_list)
jacobian = tf.expand_dims(jacobian, axis=1)
print(jacobian.shape)  # 输出 (100,1,50,1)

不过这种循环方式在y维度较大时效率不高，建议优先迁移到TF2.x使用原生方法。

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