如何从现有 Tensorflow ops 的组合中创建新的 Tensorflow op

Question

我知道如何使用tf.py_func创建一个在CPU上运行的新自定义操作。我还从 TF 指南中知道，您可以在C++中创建新的操作及其渐变

我正在寻找的不是以上任何一项。我想为 TF ops 的组合定义一个自定义渐变函数。tf.register_gradients可以与gradient_override_map一起使用，为现有操作定义自定义渐变，但是如何首先将TF ops的组合注册为新操作？

这里也提出了类似的问题，但没有答案。

Answer 1

tfe.custom_gradient是

你想使用的装饰器

Answer 2

我在此存储库中提供了三种在 Tensorflow 中定义自定义梯度的不同方法。

custom_gradient_with_py_func

在这种方法中，我们使用 tf.py_func 定义一个 tf op，并为其分配一个自定义梯度函数。

with g.gradient_override_map({"PyFunc": rnd_name}):
    return tf.py_func(func, inp, Tout, stateful=stateful, name=name)

custom_gradient_with_python：

在这种方法中，我们使用一种解决方法来定义 Tensorflow 操作组合的自定义梯度。我们覆盖标识操作的梯度。

def python_func(x_in, name=None):
    with ops.name_scope(name):
        backward_func = tf.identity(x_in) # We'll later override the gradient of identity to deflect our desired gradient function.
        forward_func = tf.subtract(2 * tf.exp(x_in), x_in) 
        return backward_func + tf.stop_gradient(forward_func - backward_func) 
def my_op(func, inp, grad, name=None, victim_op='Identity'):
    # Need to generate a unique name to avoid duplicates.
    rnd_name = 'my_gradient' + str(np.random.randint(0, 1E+8))
    tf.RegisterGradient(rnd_name)(grad)
    g = tf.get_default_graph()
    with g.gradient_override_map({victim_op: rnd_name}):
        return func(inp, name=name)

custom_gradient_with_eager：

这种方法使用 Tensorflow 1.5 中可用的 tensorflow.contrib.eager 来定义 tensorflow ops 组合的自定义梯度。

@tfe.custom_gradient
def python_func(x_in):
    def grad_func(grad):
        return grad * ((2 * tf.exp(x_in)) - 1)
    forward_func = tf.subtract(2 * tf.exp(x_in), x_in)
    return forward_func, grad_func

Answer 3

我不确定您是如何解决问题的，但上述解决方案中的名称"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。