Python 带有两个输入的Tensorflow op，返回其中一个并覆盖渐变

我试图在Tensorflow中实现一个合成梯度方案

我需要有一个op，它接受两个输入并返回其中一个（即带有虚拟变量的标识）。类似于

f（a，b）：返回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吗？

您可以在模型定义期间添加以下代码以覆盖渐变。

tf.Graph

具有实现相同功能的构造

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

---

我不确定您是如何解决问题的，但上面解决方案中的“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） out=t+tf.停止梯度（f（a，b）-t）

---

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

谢谢您的回答。我设法以这种方式获得了想要的行为。我希望能够只使用TF-ops来表达forward操作，这样代码就可以在GPU上运行了。我通过在PyTorch中实现它解决了这个问题：）无论如何，谢谢你的帮助。我很想看到有人在Tensorflow中提出直接反馈对齐的一般实现，它可能需要很多高级TF概念，我会从中学到很多东西