Python tf.custom_gradient用法及代码示例

使用自定义渐变定义函数的装饰器。

用法

tf.custom_gradient(
    f=None
)

这个装饰器允许对操作序列的梯度进行细粒度控制。由于多种原因，这可能很有用，包括为一系列操作提供更有效或数值稳定的梯度。

例如，考虑在交叉熵和对数似然计算中经常出现的以下函数：

def log1pexp(x):
  return tf.math.log(1 + tf.exp(x))

由于数值不稳定性，此函数在 x=100 处评估的梯度为 NaN。例如：

x = tf.constant(100.)
y = log1pexp(x)
dy_dx = tf.gradients(y, x) # Will be NaN when evaluated.

梯度表达式可以解析简化以提供数值稳定性：

@tf.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(upstream):
    return upstream * (1 - 1 / (1 + e))
  return tf.math.log(1 + e), grad

使用此定义，x = 100 处的梯度 dy_dx 将被正确评估为 1.0。

变量upstream 被定义为上游梯度。即来自该层的所有层或函数的梯度。上面的例子没有上游函数，因此 upstream = dy/dy = 1.0 。

假设 x_i 是前向传递中的 log1pexp x_1 = x_1(x_0) , x_2 = x_2(x_1) , ..., x_i = x_i(x_i-1) , ..., x_n = x_n(x_n-1) 。通过链式法则，我们知道 dx_n/dx_0 = dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... * dx_i/dx_i-1 * ... * dx_1/dx_0 。

在这种情况下，我们当前函数的梯度定义为 dx_i/dx_i-1 = (1 - 1 / (1 + e)) 。上游梯度 upstream 将是 dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... * dx_i+1/dx_i 。上游梯度乘以当前梯度然后向下游传递。

如果函数将多个变量作为输入，grad 函数也必须返回相同数量的变量。我们以函数z = x * y为例。

@tf.custom_gradient
def bar(x, y):
  def grad(upstream):
    dz_dx = y
    dz_dy = x
    return upstream * dz_dx, upstream * dz_dy
  z = x * y
  return z, grad
x = tf.constant(2.0, dtype=tf.float32)
y = tf.constant(3.0, dtype=tf.float32)
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  tape.watch(y)
  z = bar(x, y)
z
<tf.Tensor:shape=(), dtype=float32, numpy=6.0>
tape.gradient(z, x)
<tf.Tensor:shape=(), dtype=float32, numpy=3.0>
tape.gradient(z, y)
<tf.Tensor:shape=(), dtype=float32, numpy=2.0>

嵌套自定义渐变可能会导致不直观的结果。默认行为与n-th 阶导数不对应。例如

@tf.custom_gradient
def op(x):
  y = op1(x)
  @tf.custom_gradient
  def grad_fn(dy):
    gdy = op2(x, y, dy)
    def grad_grad_fn(ddy): # Not the 2nd order gradient of op w.r.t. x.
      return op3(x, y, dy, ddy)
    return gdy, grad_grad_fn
  return y, grad_fn

grad_grad_fn函数将计算grad_fn相对于dy的一阶梯度，用于从backward-mode梯度图生成forward-mode梯度图，但与二阶梯度不同op 相对于 x 。

相反，将嵌套的 @tf.custom_gradients 包装在另一个函数中：

@tf.custom_gradient
def op_with_fused_backprop(x):
  y, x_grad = fused_op(x)
  def first_order_gradient(dy):
    @tf.custom_gradient
    def first_order_custom(unused_x):
      def second_order_and_transpose(ddy):
        return second_order_for_x(...), gradient_wrt_dy(...)
      return x_grad, second_order_and_transpose
    return dy * first_order_custom(x)
  return y, first_order_gradient

内部 @tf.custom_gradient 装饰函数的附加参数控制最内部函数的预期返回值。

上面的示例说明了如何为不从变量读取的函数指定自定义渐变。以下示例使用变量，这些变量需要特殊处理，因为它们是 forward 函数的有效输入。

weights = tf.Variable(tf.ones([2]))  # Trainable variable weights
@tf.custom_gradient
def linear_poly(x):
  # Creating polynomial
  poly = weights[1] * x + weights[0]

  def grad_fn(dpoly, variables):
    # dy/dx = weights[1] and we need to left multiply dpoly
    grad_xs = dpoly * weights[1]  # Scalar gradient

    grad_vars = []  # To store gradients of passed variables
    assert variables is not None
    assert len(variables) == 1
    assert variables[0] is weights
    # Manually computing dy/dweights
    dy_dw = dpoly * tf.stack([x ** 1, x ** 0])
    grad_vars.append(
        tf.reduce_sum(tf.reshape(dy_dw, [2, -1]), axis=1)
    )
    return grad_xs, grad_vars
  return poly, grad_fn
x = tf.constant([1., 2., 3.])
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  poly = linear_poly(x)
poly # poly = x + 1
<tf.Tensor:shape=(3,),
  dtype=float32,
  numpy=array([2., 3., 4.], dtype=float32)>
tape.gradient(poly, x)  # conventional scalar gradient dy/dx
<tf.Tensor:shape=(3,),
  dtype=float32,
  numpy=array([1., 1., 1.], dtype=float32)>
tape.gradient(poly, weights)
<tf.Tensor:shape=(2,), dtype=float32, numpy=array([6., 3.], dtype=float32)>

上面的示例说明了可训练变量 weights 的用法。在示例中，内部 grad_fn 接受额外的 variables 输入参数，并返回额外的 grad_vars 输出。如果 forward 函数读取任何变量，则传递该额外参数。您需要计算梯度 w.r.t。每个 variables 并将其输出为 grad_vars 的列表。请注意，当 forward 函数中没有使用任何变量时，variables 的默认值设置为 None。

应该注意 tf.GradientTape 仍在观察 tf.custom_gradient 的前向传递，并将使用它观察的操作。因此，在磁带仍在观看时调用 tf.function 会导致构建梯度图。如果在 tf.function 中使用了没有注册梯度的操作，则会引发 LookupError。

用户可以插入tf.stop_gradient 来自定义此行为。这在下面的示例中得到了证明。 tf.random.shuffle 没有注册的渐变。结果 tf.stop_gradient 用于避免 LookupError 。

x = tf.constant([0.3, 0.5], dtype=tf.float32)

@tf.custom_gradient
def test_func_with_stop_grad(x):
  @tf.function
  def _inner_func():
    # Avoid exception during the forward pass
    return tf.stop_gradient(tf.random.shuffle(x))
    # return tf.random.shuffle(x)  # This will raise

  res = _inner_func()
  def grad(upstream):
    return upstream  # Arbitrarily defined custom gradient
  return res, grad

with tf.GradientTape() as g:
  g.watch(x)
  res = test_func_with_stop_grad(x)

g.gradient(res, x)

另请参阅tf.RegisterGradient，它为原始 TensorFlow 操作注册梯度函数。另一方面，tf.custom_gradient 允许对一系列操作的梯度计算进行细粒度控制。

请注意，如果修饰函数使用 Variable s，则封闭变量范围必须使用 ResourceVariable s。