tensorflow - 操作链的自定义渐变

Question

我有一个标准 TensorFlow 操作链，我需要为整个链指定一个自定义梯度。

比如说，在下面的示例中，这些操作被分组在一个 Python 函数中:'my_op'。我要做的是为“my_op”指定自定义渐变。我查看了 RegisterGradient、gradient_override_map 和 tf.Graph.create_op，但我找不到任何关于如何使用它们为一组操作定义自定义渐变而不用 C++ 重写完整操作链的简单示例。

import numpy as np
import tensorflow as tf

n = 2
m = 3
x = np.random.normal(size=(1, n))
A = tf.Variable(tf.truncated_normal(shape=(n, m), dtype=tf.float32))
b = tf.Variable(tf.zeros(shape=(1, m), dtype=tf.float32))

def my_op(a):
    return tf.add(tf.matmul(a, A), b)

x_placeholder = tf.placeholder(tf.float32,shape=[1, n])
t = my_op(tf.stop_gradient(x_placeholder))

grad = tf.gradients(t, [A])


sess = tf.Session()
sess.run(tf.initialize_all_variables())

result = sess.run(grad, feed_dict={x_placeholder: x})

print(result)

sess.close()

Answer 1

据我所知，定义自定义渐变(即对普通渐变进行一些修改)的最佳方法是在 this 之后的 tensorflow 中添加一个新的自定义操作。 .如您所见，对于输出输入的自定义操作，您可以使用 @ops.RegisterGradient("MyOp") 在 python 中定义它的梯度。