有两个输入的Tensorflow op，返回其中一个并覆盖gradien

Question

我正在尝试在Tensorflow中实现一个合成梯度方案。在

我需要一个操作，它接受两个输入并返回其中一个（即，带有伪变量的恒等式）。像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上运行，而使用py_func就不是这样了。在

我怎样才能做到这一点？我能用一下TF-op吗？在

Answer 1

我不知道你是如何解决你的问题的，但上面的解决方案中的“op_name”和“some_name”的名称不会显示在图表上。因此，您将无法使用gradient_override_map（{“op_name”：“SynthGrad”}）。在

一种可能的解决方案： 如果在forwardpass中有一个自定义的tensorflow op x=f（a，b），但希望它在backwardpass中表现为g（a，b），可以执行以下操作：

t=g（a，b） 输出=t+tf.stop_梯度（f（a，b）-t

但是，您需要在C++中定义g（a，b）作为具有名称的哑/身份操作符。稍后，您可以使用渐变覆盖图。在

Answer 2

可以在模型定义期间添加以下代码来覆盖渐变。 tf.Graph有{a1}构造来实现相同的目的

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