python - Tensorflow:如何替换或修改渐变？

Question

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

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

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

Answer 1

对于 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

以在反向 channel 中剪切渐变的层为例:

@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