如何从现有Tensorflow操作的组合创建新的Tensorflow操作系统

Question

我知道在CPU上运行的how to use tf.py_func to create a new custom op。我也从TF guide知道你可以用C ++创建一个新的op及其渐变

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

类似的问题已经被问到here没有答案。

Answer 1

tfe.custom_gradient是您要使用的装饰器

Answer 2

我在this repo中提供了三种在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操作的组合定义自定义渐变。我们覆盖了身份op的渐变。

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操作组合定义自定义渐变。

@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。