TensorFlow 自定义渐变

Question

我有一个自定义梯度计算函数，可以将传入的梯度加倍。

import tensorflow as tf

@tf.RegisterGradient("CustomSquare")
def _custom_square_grad(op, grad):
  return grad*2.0

c = tf.constant(3.)

s1 = tf.square(c)
grad1 = tf.gradients(s1, c)[0]

g = tf.get_default_graph()
with g.gradient_override_map({"Square": "CustomSquare"}):
    s2 = tf.square(c)
    grad2 = tf.gradients(s2, c)[0]

with tf.Session() as sess:
    print(sess.run([c, s1, grad1]))
    print(sess.run([c, s2, grad2]))

我得到的结果令人惊讶：

[3.0, 9.0, 6.0]
[3.0, 9.0, 2.0]

我原以为第二个结果是 [3.0, 9.0, 12.0]。我错过了什么？

谢谢。

Answer 1

简而言之，_custom_square_grad的正确版本应该是：

@tf.RegisterGradient("CustomSquare")                                             
def _custom_square_grad(op, grad):                                               
    x = op.inputs[0]                                                            
    return 2.0 * (grad * 2.0 * x)

为了理解代码，您需要知道 gradient 是如何工作的。当你定义 tf.RegisterGradient 时，它应该反向传播从输出到输入的梯度。对于tf.squre，默认的渐变函数是这样的：

# Given y = tf.square(x) => y' = 2x
grad_x = grad_y * 2.0 * x

既然你想在你自定义的渐变函数中加倍渐变，你可以简单地把它改成grad_x = 2.0 * (grad_y * 2.0 * x)。

TensorFlow 自定义渐变

TensorFlow custom gradients

backpropagation

tensorflow