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上运行，py_func也不会出现这种情况。

我怎样才能做到这一点？我可以使用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

Answer 2

我不确定您是如何设法解决问题的，但名称＆＃39; op_name＆＃39;和＆＃39; some_name＆＃39;在上面的解决方案将不会显示在图表上。所以你将无法使用gradient_override_map（{＆＃34; op_name＆＃34;：＆＃34; SynthGrad＆＃34;}）。

一种可能的解决方案：如果你在前向通道中有一个自定义张量流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。

具有两个输入的Tensorflow操作，返回两个中的一个并覆盖渐变

2 个答案: