我正在尝试为自定义TF操作定义渐变方法。我在网上找到的大部分解决方案似乎都基于gist harpone。我不愿意使用这种方法，因为它使用不会在GPU上运行的py_func。我找到了另一个使用tf.identity()的解决方案，@tf.RegisterGradient('MyCustomGradient')
def _custom_gradient(op, gradients):
x = op.inputs[0]
return(x)
def my_op(w):
return tf.pow(w,3)
var_foo = tf.Variable(5, dtype=tf.float32)
bar = my_op(var_foo)
g = tf.get_default_graph()
with g.gradient_override_map({'Identity': 'MyCustomGradient'}):
bar = tf.identity(bar)
g = tf.gradients(bar, var_foo)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print(sess.run(g))
看起来更优雅，我认为将在GPU上运行。但是，我在自定义渐变功能中访问操作的输入时遇到了一些问题。这是我的代码：
_custom_gradient()我希望op output x gradient将输入返回到op（本例中为5），但它似乎返回theListOfTasks。我的自定义my_op将具有非差异化操作，如tf.sign，我想根据输入定义自定义渐变。我究竟做错了什么？
答案 0 :(得分：2)
您的代码没有问题：
让我们先做正面传球：
var_foo = 5 - ＆gt; bar = 125 - ＆gt; tf.identity(bar) = 125
现在让我们反向传播：
tf.identity(bar)相对于其参数bar的渐变等于（根据您的定义）bar，即125。 bar相对于var_foo的渐变等于var_foo 75的平方的3倍。乘以，得到9375，这确实是代码的输出。
op.inputs[0]包含op的前向传递值。在这种情况下，identity操作的正向传递为125。