TensorFlow 自定义渐变答案

【问题标题】：TensorFlow custom gradientsTensorFlow 自定义渐变
【发布时间】：2017-10-22 15:50:42
【问题描述】：

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

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]。我错过了什么？

谢谢。

【问题讨论】：

标签： tensorflow backpropagation

【解决方案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)。

【讨论】：