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

Question

I'm trying to write a custom gradient function for 'my_op' which for the sake of the example contains just a call to tf.identity() (ideally, it could be any graph).

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()

Unfortunately I get the following error:

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>

I know that it is possible to create a custom C++ operation, but in my case I just need to write a custom gradient for a function which can be easily written in Python using standard TensorFlow operations, so I would like to avoid writing unnecessary C++ code.

Also, I'm using the upstream version of TensorFlow from GitHub.

Did you try @ops.RegisterGradient("my_op") ? You can follow the example for the python part and skip the C++ part: https://www.tensorflow.org/versions/r0.10/how_tos/adding_an_op/index.html#implement-the-gradient-in-python — Yao Zhang, Aug 09 '16 at 02:01
I think that the input to ops.RegisterGradient() is a name of a registered TensorFlow operation, it is not just the name of a Python function containing TensorFlow operations. So, somehow I need to register an operation first, right? — njk, Aug 09 '16 at 06:58
I think you are right and the code is close, but not working due to a bug here: https://github.com/tensorflow/tensorflow/issues/3710 Note that python_grad_func needs the same interface as ops.RegisterGradient https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/function.py#L349 — Yao Zhang, Aug 09 '16 at 18:42

score 3 · Accepted Answer · edited Aug 26 '16 at 11:14

Note that python_grad_func needs the same interface as ops.RegisterGradient (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/function.py#L349).

Here is the modified code example:

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()

Output:

[ 0.00669286  0.98201376  0.04742587  0.88079709  0.7310586 ]

By running your code, I got NotFoundError: Op type not registered 'my_op_2f8a34ee' — Everett You, May 29 '17 at 21:36
Running this gives me: `ValueError: Unknown keyword arguments: dict_keys(['a'])` at my_op(a). What should I do? — nikpod, Jun 12 '17 at 07:34
Found a workaround mentioned [here](https://github.com/tensorflow/tensorflow/issues/6804). Example: Define `y = my_op(a)` before `sess.run(tf.initialize_all_variables())` Solves my issue and @EverettYou issue — nikpod, Jun 12 '17 at 07:45
This didn't work for me, had to modify it as in [this gist](https://gist.github.com/gngdb/81049f9796d1292b672999f1c1f21ab7). Was able to simplify it a little though. — gngdb, Oct 26 '17 at 08:42

zfc · Answer 2 · 2017-01-13T15:46:12.933

2

The following seems work fine. Do you have any reason prefering python_grad_func instead?

@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)

edited Jan 13 '17 at 15:46

answered Sep 19 '16 at 03:50

zfc

156
2

I get `tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'fprop_da39a3ee'`. – Albert Jan 12 '17 at 10:20
Oops. sorry. Update the sample code. Let me know if it works for you. The often tricky thing here is that the function definition should be in the same Graph as its usage. – zfc Jan 13 '17 at 15:48

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

2 Answers2

Linked