Schreiben Sie benutzerdefinierte Python-basierte Verlaufsfunktion für eine Operation? (ohne C++ - Implementierung)

Question

Ich versuche, eine benutzerdefinierte Gradientenfunktion für 'my_op' zu schreiben, die für das Beispiel nur einen Aufruf von tf.identity() enthält (im Idealfall könnte es ein beliebiger Graph sein).

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

Leider bekomme ich folgende Fehlermeldung:

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> 

Ich weiß, dass es möglich ist, eine benutzerdefinierte C++ Betrieb zu schaffen, aber in meinem Fall ich brauche nur einen benutzerdefinierten Verlauf für eine Funktion zu schreiben, die kann einfach mit Standard-TensorFlow-Operationen in Python geschrieben werden, daher möchte ich vermeiden, unnötigen C++ - Code zu schreiben.

Auch ich benutze die Upstream-Version von TensorFlow von GitHub.

Answer 1

Beachten Sie, dass python_grad_func dieselbe Schnittstelle wie ops.RegisterGradient (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/function.py#L349) benötigt. Hier

ist das modifizierte Codebeispiel:

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

Ausgang:

[ 0.00669286 0.98201376 0.04742587 0.88079709 0.7310586 ] 

Answer 2

Die folgende scheint gut funktionieren. Haben Sie einen Grund, python_grad_func stattdessen zu bevorzugen?

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