Tensorflow自定义梯度：上游和输出梯度的形状

Question

由于某些操作虽然是 tf.functions，但没有定义的梯度，所以我在实现自定义损失函数时遇到了麻烦。我想我可以通过为每个操作实现 custom_gradients 来克服这个问题，但似乎我没有正确获得渐变的尺寸/形状。 遗憾的是，这也很难调试，因为除了通过 GradientTape 查看损失函数的总梯度之外，无法观察梯度在整个计算过程中到底发生了什么（至少据我所知）。 因此，想象一个带有 f 和 @tf.custom_gradient 装饰器的函数。该函数对输入执行一些操作来改变其形状。

@tf.custom_gradient
def f(input):
        output = change_shape(input)

    def grad(upstream):
        return upstream * ???

    return output, grad

例如，我如何从右侧乘以 1（类似于 tf.ones_like(some_shape)），以便上游梯度的变化符合函数 f 中张量形状的变化？ 我尝试了许多不同的方法，包括仅返回上游作为新的梯度，或者将右侧的 tf.ones_like 与输入的形状相乘（据我所知，这将是根据线性代数进行矩阵乘法的正确方法） ）。所有这些都会导致“无”梯度。

我希望我的问题是可以理解的。抱歉，没有提供任何具体代码，我现在正在旅行，缺乏资源来这样做。

谢谢

Answer 1

在 TensorFlow 中，自定义梯度允许您为自定义操作定义自己的梯度函数。实现自定义渐变时，了解上游渐变（相对于自定义操作的输出的渐变）和输出渐变（相对于自定义操作的输入的渐变）的形状非常重要。

TensorFlow 中自定义梯度函数的一般形式如下所示：

@tf.custom_gradient
def custom_op(x):
    # Forward pass logic
    
    def grad(dy):
        # Backward pass logic
        # Compute gradients with respect to x based on dy (upstream gradient)
        return dx  # Return the gradient with respect to x
    
    return y, grad

这里：

• x

  是您自定义操作的输入。

• y

  是您自定义操作的输出。

• dy

  是上游梯度，即相对于

  y

  的梯度。

现在，我们来讨论上游梯度 (

dy

) 和输出梯度 (

dx

) 的形状。

1. 上游梯度 (

   dy

   ):

   • dy

     的形状与输出

     y

     的形状相同。它表示损失相对于自定义操作的输出的梯度。

2. 输出梯度（

   dx

   ）：

   • dx

     的形状应与输入

     x

     的形状相同。它表示损失相对于自定义操作的输入的梯度。

这是一个更具体的例子：

import tensorflow as tf

@tf.custom_gradient
def custom_op(x):
    # Forward pass logic
    y = x * x
    
    def grad(dy):
        # Backward pass logic
        dx = 2 * x * dy  # Gradient with respect to x
        return dx
    
    return y, grad

# Example usage
x = tf.constant([1.0, 2.0, 3.0], dtype=tf.float32)

with tf.GradientTape() as tape:
    y = custom_op(x)

grad_x = tape.gradient(y, x)

print("Input x:", x.numpy())
print("Output y:", y.numpy())
print("Gradient with respect to x:", grad_x.numpy())

在此示例中，

custom_op

计算输入

x

的平方。自定义梯度函数计算相对于

x

的梯度，即

2 * x

。

dy

和

dx

的形状符合前面提到的规则。