具有两个输入的 Tensorflow op，返回两个输入中的一个并覆盖梯度

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 操作？

Answer 1

您可以在模型定义期间添加以下代码以覆盖梯度。tf.Graphgradient_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

Answer 2

我不确定您是如何解决问题的，但上述解决方案中的名称"op_name"和"some_name"不会显示在图表上。因此，您将无法使用gradient_override_map({"op_name"： "SynthGrad"}(。

一种可能的解决方案： 如果你在正向传递中有一个自定义的张量流 op x=f(a，b(，但你希望它在向后传递中表现为 g(a，b(，你可以做这样的事情：

t=g(a，b( out=t+tf.stop_gradient(f(a，b(-t(

但是，您需要在C++中将 g(a，b( 定义为具有名称的虚拟/标识运算符。稍后，您可以使用gradient_override_map。