使用gradient_override_map进行简单的自定义渐变

Question

我想使用一个为普通密集层创建权重的函数，它基本上就像一个初始化函数，只是它在每次新的正向传递之前“初始化”。

我的增强线性图层的流程如下所示：

input = (x, W)
W_new = g(x,W)
output = tf.matmul(x,W_new)

然而，g（x，W）不可微分，因为它涉及一些采样。幸运的是它也没有我想要学习的任何参数所以我只是尝试做前进和后退，好像我永远不会替换W. 现在我需要告诉自动区分不通过g（）反向传播。我这样做：

W_new = tf.stop_gradient(g(x,W))

不幸的是，这不起作用，因为它抱怨不匹配的形状。 工作如下：

input = (x, W)
W_new = W + tf.stop_gradient(g(x,W) - W)
output = tf.matmul(x,W_new)

如此处所示：https://stackoverflow.com/a/36480182

现在前进似乎没问题，但我不知道如何覆盖后向传球的渐变。我知道，我必须使用：gradient_override_map，但无法将我见过的应用程序转移到我的特定用例（我对TF仍然很新）。 但是，我不知道如何做到这一点，如果没有更简单的方法。我假设在给定模型中的第一个前向传递中必须完成类似的操作，其中所有权重都被初始化，而我们也不必通过init函数反向传播。

非常感谢任何帮助！

Answer 1

嘿@jhj我也遇到了同样的问题，幸运的是我找到了gist。希望这会有所帮助：）

样本工作 -

将tensorflow导入为tf

来自tensorflow.python.framework import ops

将numpy导入为np

定义自定义py_func，它还将一个grad op作为参数：

def py_func（func，inp，Tout，stateful = True，name = None，grad = None）：

# Need to generate a unique name to avoid duplicates:
rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

tf.RegisterGradient(rnd_name)(grad)  # see _MySquareGrad for grad example
g = tf.get_default_graph()
with g.gradient_override_map({"PyFunc": rnd_name, "PyFuncStateless": rnd_name}):
    return tf.py_func(func, inp, Tout, stateful=stateful, name=name)

使用np.square而不是tf.square：

来定义自定义方形函数

def mysquare（x，name = None）：

with ops.name_scope(name, "Mysquare", [x]) as name:
    sqr_x = py_func(np.square,
                    [x],
                    [tf.float32],
                    name=name,
                    grad=_MySquareGrad)  # <-- here's the call to the gradient
    return sqr_x[0]

实际渐变：

def _MySquareGrad（op，grad）：

x = op.inputs[0]
return grad * 20 * x  # add a "small" error just to see the difference:

使用tf.Session（）作为sess：

x = tf.constant([1., 2.])
y = mysquare(x)
tf.global_variables_initializer().run()

print(x.eval(), y.eval(), tf.gradients(y, x)[0].eval())