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: tensorflow.org/versions/r0.10/how_tos/adding_an_op/… — Yao Zhang
– Yao Zhang, Commented Aug 9, 2016 at 2: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
– njk, Commented Aug 9, 2016 at 6:58
I think you are right and the code is close, but not working due to a bug here: github.com/tensorflow/tensorflow/issues/3710 Note that python_grad_func needs the same interface as ops.RegisterGradient github.com/tensorflow/tensorflow/blob/master/tensorflow/python/… — Yao Zhang
– Yao Zhang, Commented Aug 9, 2016 at 18:42

Ohad Eytan · Accepted Answer · 2016-08-26 11:14:19Z

3

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 ]

edited Aug 26, 2016 at 11:14

Ohad Eytan

8,5341 gold badge24 silver badges32 bronze badges

answered Aug 10, 2016 at 0:22

Yao Zhang

5,8111 gold badge18 silver badges22 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Everett You Over a year ago

By running your code, I got NotFoundError: Op type not registered 'my_op_2f8a34ee'

nikpod Over a year ago

Running this gives me: ValueError: Unknown keyword arguments: dict_keys(['a']) at my_op(a). What should I do?

nikpod Over a year ago

Found a workaround mentioned here. Example: Define y = my_op(a) before sess.run(tf.initialize_all_variables()) Solves my issue and @EverettYou issue

gngdb Over a year ago

This didn't work for me, had to modify it as in this gist. Was able to simplify it a little though.

zfc · Accepted Answer · 2017-01-13 15:46:12Z

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, 2017 at 15:46

answered Sep 19, 2016 at 3:50

zfc

1562 bronze badges

2 Comments

Albert Over a year ago

I get tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'fprop_da39a3ee'.

zfc Over a year ago

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.

Collectives™ on Stack Overflow

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

2 Answers 2

4 Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

4 Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related