TensorFlow中能否单次图运行高效计算逐样本梯度？

当然可以实现！其实你要的就是逐样本梯度（per-sample gradients），TensorFlow里有现成的方式帮你在单次图运行里完成这个操作，完全不用循环逐个计算样本梯度。

针对你举的简单示例的实现 #

你不用自定义tf.multiple_gradients，用原生的tf.gradients就能直接搞定，注意维度匹配就行：

import tensorflow as tf

x = tf.placeholder(tf.float32, shape=[100])
f = tf.square(x)
# 计算f对x的梯度，默认就是逐元素返回，对应每个x_i的f'(x_i)
f_grad = tf.gradients(f, x)[0]

运行这个图时，f_grad会直接输出[2*x[0], 2*x[1], ..., 2*x[99]]，完全是向量化的单次计算，效率拉满。

适配黑箱随机变分推断（BBSVI）的场景 #

在BBSVI里，你的损失通常是一批样本的期望，这时候需要每个样本对应的梯度来更新变分参数，这里分两种版本给你说：

TensorFlow 2.x 版本（更直观）

用tf.GradientTape可以直接拿到逐样本梯度：

import tensorflow as tf

x = tf.Variable(tf.random.normal([100]))
with tf.GradientTape() as tape:
    f = tf.square(x)
# 直接得到逐元素的梯度，也就是每个样本对应的导数
f_grad = tape.gradient(f, x)

TensorFlow 1.x 版本

如果处理的是批量样本（比如输入形状是[batch_size, feature_dim]），要把损失构造成逐样本的张量形式（不是求和后的标量），这样tf.gradients会自动返回每个样本的梯度：

# 假设x是形状为[100, 10]的批量输入
x = tf.placeholder(tf.float32, shape=[100, 10])
# 每个样本的损失，形状保持为[100]
per_sample_loss = tf.reduce_sum(tf.square(x), axis=1)
# 计算每个样本损失对x的梯度，结果形状为[100,10]，对应每个样本的梯度
per_sample_grads = tf.gradients(per_sample_loss, x)[0]

核心要点 #

• 关键是让损失函数保持逐样本的张量形态，而不是求和压缩成标量，这样TensorFlow的自动微分机制会自动为每个样本计算对应的梯度，全程单次图运行，效率很高。
• 不需要手动写循环遍历样本，原生API就支持向量化的逐样本梯度计算。

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