为操作编写基于 Python 的自定义梯度函数？（没有 C++ 实现）

Question

我正在尝试为 'my_op' 编写一个自定义梯度函数，为了示例，它只包含对 tf.identity() 的调用（理想情况下，它可以是任何图形）。

import tensorflow as tf
from tensorflow.python.framework import function


def my_op_grad(x):
    return [tf.sigmoid(x)]


@function.Defun(a=tf.float32, python_grad_func=my_op_grad)
def my_op(a):
    return tf.identity(a)


a = tf.Variable(tf.constant([5., 4., 3., 2., 1.], dtype=tf.float32))

sess = tf.Session()
sess.run(tf.initialize_all_variables())

grad = tf.gradients(my_op(a), [a])[0]

result = sess.run(grad)

print(result)

sess.close()

不幸的是，我收到以下错误：

Traceback (most recent call last):
  File "custom_op.py", line 19, in <module>
    grad = tf.gradients(my_op(a), [a])[0]
  File "/Users/njk/tfm/lib/python3.5/site-packages/tensorflow/python/framework/function.py", line 528, in __call__
    return call_function(self._definition, *args, **kwargs)
  File "/Users/njk/tfm/lib/python3.5/site-packages/tensorflow/python/framework/function.py", line 267, in call_function
    compute_shapes=False)
  File "/Users/njk/tfm/lib/python3.5/site-packages/tensorflow/python/framework/ops.py", line 2285, in create_op
    raise TypeError("Input #%d is not a tensor: %s" % (idx, a))
TypeError: Input #0 is not a tensor: <tensorflow.python.ops.variables.Variable object at 0x1080d2710>

我知道可以创建自定义 C++ 操作，但在我的例子中，我只需要为函数编写自定义渐变，可以使用标准 TensorFlow 操作轻松地在 Python 中编写，所以我想避免编写不必要的 C++ 代码。

此外，我正在使用来自 GitHub 的上游版本的 TensorFlow。

Answer 1

请注意 python_grad_func 需要与 ops.RegisterGradient (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/function.py#L349) 相同的接口。

修改后的代码示例如下：

def my_op_grad(op, grad): ### instead of my_op_grad(x)                                                  
    return tf.sigmoid(op.inputs[0])                                              

@function.Defun(a=tf.float32, python_grad_func=my_op_grad)                       
def my_op(a):                                                                    
    return tf.identity(a)                                                        

def main(unused_argv):                                                           

  a = tf.Variable(tf.constant([-5., 4., -3., 2., 1.], dtype=tf.float32))         
  sess = tf.Session()                                                            
  sess.run(tf.initialize_all_variables())                                        

  a = tf.identity(a) #workaround for bug github.com/tensorflow/tensorflow/issues/3710

  grad = tf.gradients(my_op(a), [a])[0]                                          
  result = sess.run(grad)                                                        

  print(result)                                                                  

  sess.close()

输出：

[ 0.00669286  0.98201376  0.04742587  0.88079709  0.7310586 ]

Answer 2

以下似乎工作正常。您有什么理由更喜欢 python_grad_func 吗？

@tf.function.Defun(tf.float32, tf.float32)
def bprop(x, dy):
  return tf.sigmoid(x)

@tf.function.Defun(tf.float32, grad_func=bprop)
def fprop(x):
  return x  # identity

a = tf.Variable(tf.constant([-5., 4., -3., 2., 1.], dtype=tf.float32))
grad = tf.gradients(fprop(a), [a])                                         

with tf.Session() as sess:                                                             
  sess.run(tf.initialize_all_variables())
  result = sess.run(grad)                                                        

print(result)

为操作编写基于 Python 的自定义梯度函数？（没有 C++ 实现）

Write Custom Python-Based Gradient Function for an Operation? (without C++ Implementation)

python

gradient-descent

tensorflow

为操作编写基于 Python 的自定义梯度函数？ （没有 C++ 实现）

Write Custom Python-Based Gradient Function for an Operation? (without C++ Implementation)

python

gradient-descent

tensorflow

为操作编写基于 Python 的自定义梯度函数？（没有 C++ 实现）