如何为多输出TensorFlow函数编写自定义梯度？代码异常排查

核心问题点 #

1. 导数计算的数学错误：
   tan(z)对z的导数是1/cos²(z)（其中z = x@w + b），但原代码错误地使用y = tan(z)计算cos(y)，导致梯度完全偏离正确值。
2. 自定义梯度的签名错误：
   tf.custom_gradient要求梯度函数返回两部分：
   • 第一部分是函数输入（此处为x）的梯度
   • 第二部分是一个函数，用于接收需要求梯度的变量列表并返回对应梯度列表
     原代码直接将变量梯度列表与输入梯度一起返回，不符合API规范。
3. 冗余的嵌套tf.function：
   自定义梯度内部嵌套@tf.function会干扰TensorFlow的梯度追踪流程，属于不必要的封装。

修正后的代码 #

import tensorflow as tf

#-----------------------------------------------------------------
# Gradient test example 
#-----------------------------------------------------------------
w = tf.Variable([[1.,2.],[2.,3.],[3.,4.]], name='w') # shape (3,2)
b = tf.Variable([1.0,2.0], name='b') # shape (2,)
x = tf.constant([[1., 2., 3.],[1.,1.,1.]]) # shape (2,3)

@tf.custom_gradient
def custom_op(x):
    # 前向计算：先计算z = x@w + b，再求tan(z)
    z = x @ w + b
    y = tf.math.tan(z)
    
    def grad(upstream):
        # 计算tan(z)对z的导数：1/cos²(z)
        yp = 1.0 / tf.square(tf.math.cos(z))
        # 计算输入x的梯度
        dydx = (upstream * yp) @ tf.transpose(w)
        
        # 定义计算变量梯度的函数
        def grad_vars(variables):
            var_grads = []
            for var in variables:
                if var is w:
                    # 对w的梯度：x^T @ (upstream * yp)
                    var_grads.append(tf.transpose(x) @ (upstream * yp))
                elif var is b:
                    # 对b的梯度：按batch维度求和
                    var_grads.append(tf.reduce_sum(upstream * yp, axis=0))
                else:
                    var_grads.append(None)
            return var_grads
        
        return dydx, grad_vars
    
    return y, grad

#-----------------------------------------------
# 前向传播与梯度计算
#-----------------------------------------------
with tf.GradientTape(persistent=True) as tape:
    tape.watch(x)
    y = tf.math.tan(x @ w + b )
    y2 = custom_op(x)
    loss = tf.reduce_mean(y**2)    
    loss2 = tf.reduce_mean(y2**2)    

# 计算损失对各变量/输入的梯度
my_vars = {'w': w,'b': b}
dldx = tape.gradient(loss,x)
dldwb = tape.gradient(loss, my_vars)
dldx_2 = tape.gradient(loss2,x)
dldwb_2 = tape.gradient(loss2, my_vars)

# 计算y对各变量/输入的梯度
dydx = tape.gradient(y,x)  
dydw = tape.gradient(y,w)  
dydb = tape.gradient(y,b)
dydx_2 = tape.gradient(y2,x)
dydw_2 = tape.gradient(y2,w)
dydb_2 = tape.gradient(y2,b)

# 打印验证结果
print('w :', w.numpy())
print('b :', b.numpy())
print('x :', x.numpy())    
print('y 与 y2 是否一致:', tf.reduce_all(tf.equal(y, y2)).numpy())
print('loss 与 loss2 是否一致:', tf.equal(loss, loss2).numpy())
print('dldx 与 dldx_2 是否一致:', tf.reduce_all(tf.abs(dldx - dldx_2) < 1e-6).numpy())
print('dldw 与 dldw_2 是否一致:', tf.reduce_all(tf.abs(dldwb['w'] - dldwb_2['w']) < 1e-6).numpy())
print('dldb 与 dldb_2 是否一致:', tf.reduce_all(tf.abs(dldwb['b'] - dldwb_2['b']) < 1e-6).numpy())
print('dydx 与 dydx_2 是否一致:', tf.reduce_all(tf.abs(dydx - dydx_2) < 1e-6).numpy())
print('dydw 与 dydw_2 是否一致:', tf.reduce_all(tf.abs(dydw - dydw_2) < 1e-6).numpy())
print('dydb 与 dydb_2 是否一致:', tf.reduce_all(tf.abs(dydb - dydb_2) < 1e-6).numpy())

验证结果 #

修正后，所有梯度结果会与常规方式计算的结果完全一致（误差小于1e-6，处于浮点计算精度范围内）。

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