tensorflow - 创建自定义图层

Question

我正在尝试在 tensorflow 中创建一个标志层。因为 Sign 的梯度为零，所以我想创建一个这样的函数：

但我不知道应该在哪里写backward函数？

标签： tensorflowkerascustomizationkeras-layertf.keras

Answer 1

一般来说，这可以通过tf.custom_gradient. 这允许您编写一个前向函数以及一个自定义梯度函数，该函数取决于“传入”梯度（从模型更下方的层）以及前向函数。

有了它，你可以创建一个这样的函数：

@tf.custom_gradient
def sign_with_grad(x):
    output = tf.sign(x)
    def grad_fn(dy):
        check_range = tf.where(tf.less_equal(tf.abs(x), 1.), 1., 0.)
        return dy*check_range
    return output, grad_fn

在这里，我们编写了一个自定义梯度，它简单地穿过上面的层——除了它将输入超出范围 [-1, 1] 的所有元素归零。我们返回正向函数的结果（符号）以及梯度函数本身。装饰者负责处理其余部分。

请注意，我没有检查代码是否运行——如果没有，请告诉我！特别是，less_equaland 或where检查可能需要显式广播——例如使用tf.ones_like(x)代替1.（对于 0 相同）。