TensorFlow 2.2向量/张量输入输出自定义层的自定义梯度实现问询

1. custom_grad返回值的数学含义 #

你写的custom_grad函数的输入dy是损失函数对custom_op输出result的梯度，数学表达式为$\partial L / \partial y$，形状和result完全一致。
custom_grad的返回值需要和custom_op的输入参数一一对应，每个返回值分别对应损失函数对对应输入参数的梯度，形状必须和对应输入完全相同。你示例里custom_op输入是A和W两个参数，因此custom_grad需要返回两个梯度，分别是$\partial L / \partial A$和$\partial L / \partial W$。
TensorFlow的反向传播采用向量-雅可比乘积(VJP)机制，你不需要显式构造完整的雅可比矩阵，直接返回每个输入对应的损失梯度即可，框架会自动完成后续的链式传播计算。

2. 示例的梯度实现 #

针对你给出的y = tf.matmul(A, W)场景，我们先明确各张量形状：

• A形状：(样本数N, 输入维度in_dim)
• W形状：(in_dim, 输出维度out_dim)
• 输出y形状：(N, out_dim)，对应dy形状也为(N, out_dim)
  根据矩阵求导的链式法则，两个梯度的计算逻辑为：
• 对输入A的梯度：$\partial L / \partial A = dy \cdot W^T$，形状为(N, in_dim)，和A形状一致
• 对权重W的梯度：$\partial L / \partial W = A^T \cdot dy$，形状为(in_dim, out_dim)，和W形状一致
  补全后的代码如下：

import tensorflow as tf

@tf.custom_gradient
def custom_op(A,W):
    result = tf.matmul(A, W)
    def custom_grad(dy):
        grad_A = tf.matmul(dy, W, transpose_b=True)
        grad_W = tf.matmul(A, dy, transpose_a=True)
        return grad_A, grad_W
    return result, custom_grad

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, out_dim):
        super().__init__()
        self.out_dim = out_dim
    
    def build(self, input_shape):
        # 初始化权重W，input_shape[-1]是输入维度
        self.W = self.add_weight(shape=(input_shape[-1], self.out_dim),
                                 initializer='random_normal',
                                 trainable=True)
    
    def call(self, A):
        return custom_op(A, self.W)

3. 梯度正确性验证 #

你可以通过tf.GradientTape对比自定义梯度和原生实现的梯度是否一致，确认逻辑正确：

# 测试数据
A = tf.random.normal((32, 16)) # 32个样本，输入维度16
layer = CustomLayer(8) # 输出维度8

# 计算自定义层的梯度
with tf.GradientTape() as tape:
    y_custom = layer(A)
    loss = tf.reduce_sum(y_custom)
grads_custom = tape.gradient(loss, layer.W)

# 计算原生matmul的梯度
with tf.GradientTape() as tape:
    y_native = tf.matmul(A, layer.W)
    loss_native = tf.reduce_sum(y_native)
grads_native = tape.gradient(loss_native, layer.W)

# 验证梯度一致
print(tf.reduce_all(grads_custom == grads_native)) # 输出tf.Tensor(True, shape=(), dtype=bool)

4. 通用实现逻辑 #

所有自定义梯度场景都可以遵循以下流程开发：

• 确认自定义op的所有输入参数数量，custom_grad返回值数量和输入参数数量严格一致
• 每个返回值的形状和对应输入的形状完全相同
• 不需要构造完整雅可比矩阵，直接计算损失对每个输入的梯度即可

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