Jak zarejestrować niestandardowy gradient dla operacji składającej się z operacji Tf

Question

Dokładniej mam prosty fprop, który jest składem operacji TF. Chcę zastąpić obliczenia gradientu tensorflow z własnej metody gradientu przy użyciu RegisterGradient.

Co jest nie tak z tym kodem?

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

@ops.RegisterGradient("MyopGrad") 
def frop_grad(op, grad): 
    x = op.inputs[0] 
    return 0 * x # zero out to see the difference: 

def fprop(x): 
    x = tf.sqrt(x) 
    out = tf.maximum(x, .2) 
    return out 

a = tf.Variable(tf.constant([5., 4., 3., 2., 1.], dtype=tf.float32)) 
h = fprop(a) 
h = tf.identity(h, name="Myop") 
grad = tf.gradients(h, a) 

g = tf.get_default_graph() 
with g.gradient_override_map({'Myop': 'MyopGrad'}): 
    with tf.Session() as sess: 
     sess.run(tf.initialize_all_variables()) 
     result = sess.run(grad) 

print(result[0]) 

chcę zobaczyć wszystkie zera w druku, ale zamiast ja dostaję:

[ 0.2236068 0.25000003 0.28867513 0.35355341 0.5  ] 

Answer 1

Trzeba zdefiniować op zakresem with g.gradient_override_map({'Myop': 'MyopGrad'})

Ponadto, należy do mapa Identity zamiast nazwy Myop do nowego gradientu.

Oto pełny kod:

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

@ops.RegisterGradient("MyopGrad") 
def frop_grad(op, grad): 
    x = op.inputs[0] 
    return 0 * x # zero out to see the difference: 

def fprop(x): 
    x = tf.sqrt(x) 
    out = tf.maximum(x, .2) 
    return out 

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

g = tf.get_default_graph() 
with g.gradient_override_map({'Identity': 'MyopGrad'}): 
    h = tf.identity(h, name="Myop") 
    grad = tf.gradients(h, a) 

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

print(result[0]) 

wyjściowa:

[ 0. 0. 0. 0. 0.]