如何註冊自定義梯度爲TF操作組成

Question

更具體我有一個簡單fprop即TF操作的複合操作。 我想用我自己的梯度法使用RegisterGradient來重寫張量梯度計算。

有什麼不對的代碼？

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

我想看看在打印全部爲零，而是我得到：

[ 0.2236068 0.25000003 0.28867513 0.35355341 0.5  ] 

Answer 1

您需要的with g.gradient_override_map({'Myop': 'MyopGrad'})

而且範圍內定義運算，你需要地圖Identity而不是名稱Myop到您的新梯度。

下面是完整的代碼：

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

輸出：

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