在Keras中定义自定义渐变时出错

Question

我一直在尝试用自定义离散梯度定义Keras中的自定义层，因为激活函数是离散的

该层如下所示：

class DiffLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(DiffLayer, self).__init__()

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(15, 1),
            initializer="random_normal",
            trainable=True,
        )

        self.b = self.add_weight(
            shape=(1, 1), initializer="random_normal", trainable=True
        )

    def call(self, x):
        z = tf.matmul(Flatten()(x), self.w) + self.b
        a = custom_op(z)
        self.a = a
        if K.greater(a,0.5):
            return x-1
        else: 
            return x

以及custom_op函数：

@tf.custom_gradient
def custom_op(x):
    a = 1. / (1. + K.exp(-x))
    def custom_grad(dy):
        if K.greater(a, 0.5):
          grad = K.exp(x)
        else:
          grad = 0
        return grad
    return a, custom_grad

我遵循了this post的教程，但当我尝试适应我正在使用的网络时，我得到以下警告：

WARNING:tensorflow:Gradients do not exist for variables ['diff_layer_10/Variable:0', 'diff_layer_10/Variable:0'] when minimizing the loss.

我的猜测是，由于定义的方式，Keras没有检测到定义的梯度，但我想不出另一种定义它的方式

是这样，还是我的代码中遗漏了什么

编辑

正如其中一条评论所建议的，我将进一步解释我试图做的事情。我希望a是一个决定输入数据发生什么变化的参数。如果a大于0.5，那么我希望输入数据减去1，否则层应该返回输入数据

我不知道在凯拉斯是否可以做到这一点