如何用tf.GradientTape实现多输出神经网络的自动微分？

针对双输入（x、y）双输出（f1、f2）的神经网络，要获取df1/dx、df1/dy、df2/dx、df2/dy四个偏导，只需在tf.GradientTape上下文内计算所有输出，再分别对每个输出调用梯度计算即可，具体实现如下：

代码示例 #

import tensorflow as tf

# 示例双输出神经网络（替换为你的实际模型）
def dual_output_model(x, y):
    f1 = x ** 2 + 3 * y
    f2 = tf.sin(x) + tf.exp(y)
    return f1, f2

# 定义可训练的输入张量
x = tf.Variable(1.0)
y = tf.Variable(2.0)

# 启用GradientTape并设置persistent=True（支持多次梯度计算）
with tf.GradientTape(persistent=True) as tape:
    # 在tape上下文内计算两个输出，确保运算被追踪
    f1, f2 = dual_output_model(x, y)

# 分别计算每个输出对输入的偏导
df1_dx = tape.gradient(f1, x)
df1_dy = tape.gradient(f1, y)
df2_dx = tape.gradient(f2, x)
df2_dy = tape.gradient(f2, y)

# 释放tape资源
del tape

# 输出结果
print(f"df1/dx: {df1_dx.numpy()}")
print(f"df1/dy: {df1_dy.numpy()}")
print(f"df2/dx: {df2_dx.numpy()}")
print(f"df2/dy: {df2_dy.numpy()}")

关键说明 #

• 追踪运算：必须在tf.GradientTape的代码块内计算所有需要求导的输出（f1、f2），否则tape无法记录运算路径，无法计算梯度。
• persistent=True：默认情况下，tape只能调用一次gradient()方法，设置该参数后允许多次调用，满足对两个输出分别求导的需求。
• 资源释放：使用完tape后手动del tape，避免内存泄漏。

也可以采用和单输出场景类似的批量求导方式，代码更简洁：

with tf.GradientTape(persistent=True) as tape:
    f1, f2 = dual_output_model(x, y)

# 一次性获取f1对x、y的偏导
df1_dx, df1_dy = tape.gradient(f1, [x, y])
# 一次性获取f2对x、y的偏导
df2_dx, df2_dy = tape.gradient(f2, [x, y])

del tape

内容的提问来源于stack exchange，提问作者sing a song