Tensorflow：如何替换或修改渐变？

Posted 2023-02-23

【中文标题】Tensorflow：如何替换或修改渐变？

【英文标题】：Tensorflow: How to replace or modify gradient?

【发布时间】：2017-10-05 23:18:26

【问题描述】：

我想替换或修改 tensorflow 中操作或图形部分的梯度。如果我可以在计算中使用现有的梯度，那将是理想的。

在某些方面，这与 tf.stop_gradient() 所做的相反：我想要一个仅在计算梯度时使用的计算，而不是添加一个在计算梯度时被忽略的计算。

一个简单的例子是通过将梯度乘以一个常数来简单地缩放梯度（但不会将前向计算乘以一个常数）。另一个例子是将渐变剪裁到给定范围内。

【参考方案1】：

对于 TensorFlow 2，您应该使用 tf.custom_gradient 装饰器，如下所示：

@tf.custom_gradient
def func(x):
    f = # calculate forward pass
    def grad(dy):
        gradient = # calculate custom gradient of func
        return dy * gradient
    return f, grad

请注意，您必须将梯度乘以上游梯度。不过要小心！

如果您在创建 Keras 函数模型时将其作为函数调用并使用 tf.GradientTape，则仍然会发生自动微分，并且您的自定义渐变将被忽略。

相反，您必须将您的函数放入一个层中：

class func_layer(tf.keras.layers.Layer):
    def __init__(self):
        super(func_layer, self).__init__()

    def call(self, x):
        return func(x)

现在，当您向函数模型添加 func_layer 时，将适当地计算后向传递。

【参考方案2】：

对于 TensorFlow 1.7 和 TensorFlow 2.0，请查看编辑打击。

---

首先定义你的自定义渐变：

@tf.RegisterGradient("CustomGrad")
def _const_mul_grad(unused_op, grad):
  return 5.0 * grad

由于您不希望在前向传播中发生任何事情，因此用您的新梯度覆盖恒等运算的梯度：

g = tf.get_default_graph()
with g.gradient_override_map("Identity": "CustomGrad"):
  output = tf.identity(input, name="Identity")

这是一个工作示例，它使用相同的方法在向后传递中剪切渐变并且在向前传递中不执行任何操作：

import tensorflow as tf

@tf.RegisterGradient("CustomClipGrad")
def _clip_grad(unused_op, grad):
  return tf.clip_by_value(grad, -0.1, 0.1)

input = tf.Variable([3.0], dtype=tf.float32)

g = tf.get_default_graph()
with g.gradient_override_map("Identity": "CustomClipGrad"):
  output_clip = tf.identity(input, name="Identity")
grad_clip = tf.gradients(output_clip, input)

# output without gradient clipping in the backwards pass for comparison:
output = tf.identity(input)
grad = tf.gradients(output, input)

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print("with clipping:", sess.run(grad_clip)[0])
  print("without clipping:", sess.run(grad)[0])

---

为 TensorFlow 1.7 和 TensorFlow 2.0 编辑

从 1.7 开始，有一种新的方法可以用更短的语法重新定义梯度，这也适用于 Tensorflow 2.0。它还允许同时重新定义多个操作的梯度。以下是上面的示例，为 TensorFlow 1.7 和 TensorFlow 2.0 重写：

在反向传播中缩放渐变的层：

@tf.custom_gradient
def scale_grad_layer(x):
  def grad(dy):
    return 5.0 * dy
  return tf.identity(x), grad

以在反向通道中剪切渐变的层为例：

@tf.custom_gradient
def clip_grad_layer(x):
  def grad(dy):
    return tf.clip_by_value(dy, -0.1, 0.1)
  return tf.identity(x), grad

【讨论】：

这是否会同时修改链中的后续渐变？

@KevinP 例如，对于裁剪：在恒等运算的反向传递过程中，渐变只会被裁剪 1 次。但是链中的所有先前层都会受到影响，因为每一层都使用其下一层的梯度进行反向传播。但是之前的图层他们自己不会再次剪辑。

谢谢。整个反向传播与前向传播使这个问题比预期的更令人困惑。我的意思是后面的反向传播梯度链。

grad 能否接受其他参数，例如，一些用于减少计算的中间变量？

@HuangYuheng grad 不能接受额外的参数，但可以使用 tensorflow 变量，而这些变量又可以在 forward pass 中改变。

【参考方案3】：

对于当前的 TensorFlow r1.13，请使用 tf.custom_gradient。

装饰函数（输入参数是一个列表x）应该返回

前传的结果，以及 一个函数，它返回一个渐变列表，x 中的每个元素对应一个。

这是一个只有一个变量的例子：

@tf.custom_gradient
def non_differentiable(x):
    f = tf.cast(x > 0, tf.float32)
    def grad(dy):
        return tf.math.maximum(0., 1 - tf.abs(x))
    return f, grad

一对二：

@tf.custom_gradient
def non_differentiable2(x0, x1):
    f = x0 * tf.cast(x1 > 0, tf.float32)
    def grad(dy):
        df_dx0 = tf.cast(x1 > 0, tf.float32)
        return dy*df_dx0, tf.zeros_like(dy)
    return f, grad

【讨论】：

嗨啦啦伙伴，谢谢你的回答。你知道如何改变 relu 函数的梯度吗？

【参考方案4】：

假设前向计算为

y = f(x)

你希望它像反向传播

y = b(x)

一个简单的技巧是：

y = b(x) + tf.stop_gradient(f(x) - b(x))

【参考方案5】：

最常用的方法是使用 https://www.tensorflow.org/api_docs/python/tf/RegisterGradient

下面，我实现了反向传播渐变剪裁，它可以与matmul 一起使用，如此处所示，或任何其他操作：

import tensorflow as tf
import numpy as np

# from https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):

    # Need to generate a unique name to avoid duplicates:
    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

    tf.RegisterGradient(rnd_name)(grad)
    g = tf.get_default_graph()
    with g.gradient_override_map("PyFunc": rnd_name):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)

def clip_grad(x, clip_value, name=None):
    """"
    scales backpropagated gradient so that
    its L2 norm is no more than `clip_value`
    """
    with tf.name_scope(name, "ClipGrad", [x]) as name:
        return py_func(lambda x : x,
                        [x],
                        [tf.float32],
                        name=name,
                        grad=lambda op, g : tf.clip_by_norm(g, clip_value))[0]

---

示例用法：

with tf.Session() as sess:
    x = tf.constant([[1., 2.], [3., 4.]])
    y = tf.constant([[1., 2.], [3., 4.]])

    print('without clipping')
    z = tf.matmul(x, y)
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval())

    print('with clipping')
    z = tf.matmul(clip_grad(x, 1.0), clip_grad(y, 0.5))
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval())

    print('with clipping between matmuls')
    z = tf.matmul(clip_grad(tf.matmul(x, y), 1.0), y)
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval())

---

输出：

without clipping
[[ 3.  7.]
 [ 3.  7.]]
with clipping
[[ 0.278543   0.6499337]
 [ 0.278543   0.6499337]]
with clipping between matmuls
[[ 1.57841039  3.43536377]
 [ 1.57841039  3.43536377]]

【讨论】：

MaxB：谢谢！这看起来很有用。我不确定如何在 python 中定义一个新的操作...它只是一个带有装饰器的函数吗？你能做一个带有剪裁渐变的 matmult 的完整示例吗？

@AlexI 这并不容易，但它是可行的：***.com/questions/37924071/… 如果您只想剪切渐变，我建议您定义一个除了剪切渐变之外什么都不做的“身份操作”。另请参阅tensorflow.org/extend/…

@AlexI 我实现了实际的反向传播渐变裁剪。见编辑

【参考方案6】：

使用optimizer.compute_gradients 或tf.gradient 获取原始渐变 然后为所欲为 最后，使用optimizer.apply_gradients

我从 github 找到了一个example

【讨论】：

谢谢，这很有趣。我认为它取代了完整的（端到端）渐变，并且仅用于优化器。我想替换单个操作的梯度，同时让其他操作的梯度以正常方式传播；我不一定知道如何处理端到端渐变。一个例子是有一个 tf.matmult() ，其中前向计算正常完成，但梯度是 clip(grad, min, max) ，其中 grad 是原始梯度，并在更大的图中使用。

看看compute_gradients，它返回一个(gradient, variable)对的列表，所以我认为你只能修改你想要的“the”渐变，比如this，找到你想要的var