帶有兩個輸入的Tensorflow op，返回其中一個並覆蓋梯度

Question

我試圖在Tensorflow中實現合成梯度方案。

我需要一個op，它接受兩個輸入並返回其中的一個（即帶有虛擬變量的身份）。像f(a, b): return a

東西，我需要這個，因爲當時我想這取決於b

@ops.RegisterGradient("SynthGrad") 
def _SynthGrad(op, grad): 
dim1 = tf.shape(op.inputs[1])[1] 
dim2 = tf.shape(op.inputs[0])[1] 
B = tf.random_normal([dim1, dim2]) 
synth_grad = tf.matmul(op.inputs[1], B) 
return synth_grad 

編寫Python函數，並使用tf.py_func將是一個爛攤子一個公式來覆蓋梯度。另外，我想讓這個代碼也運行在GPU上，並且不會出現這種情況。

我該如何做到這一點？有沒有可以使用的TF操作？

Answer 1

您可以在模型定義期間添加以下代碼以覆蓋漸變。 tf.Graph有gradient_override_map結構來達到同樣的

g = tf.get_default_graph() 
...model, definiton, input other op etc 

# gradient overrring map construct with the function `f` in your case 
with g.gradient_override_map({"op_name": "SynthGrad"}): 
    f_out = f(op_in_1, op_in_2, name="op_name") 

... 
# code related to custom function and custom gradient from your question 
def f(a, b, name='some_name'): 
    ... some stuffs 
    return a  

@tf.RegisterGradient("SynthGrad") 
def _SynthGrad(op, grad): 
    dim1 = tf.shape(op.inputs[1])[1] 
    dim2 = tf.shape(op.inputs[0])[1] 
    B = tf.random_normal([dim1, dim2]) 
    synth_grad = tf.matmul(op.inputs[1], B) 
    return synth_grad