python之在 TensorFlow 中实现重步函数

我想在 TensorFlow 中创建 heaviside step 函数。由于 Heaviside 函数不可微分，我还需要选择导数近似并定义自定义梯度，因此完整实现如下所示:

import tensorflow as tf 
 
 
@tf.RegisterGradient("HeavisideGrad") 
def _heaviside_grad(unused_op: tf.Operation, grad: tf.Tensor): 
    x = unused_op.inputs[0] 
    # During backpropagation heaviside behaves like sigmoid 
    return tf.sigmoid(x) * (1 - tf.sigmoid(x)) * grad 
 
 
def heaviside(x: tf.Tensor, g: tf.Graph = tf.get_default_graph()): 
    custom_grads = { 
        "Sign": "HeavisideGrad" 
    } 
    with g.gradient_override_map(custom_grads): 
        # TODO: heaviside(0) currently returns 0. We need heaviside(0) = 1 
        sign = tf.sign(x) 
        # tf.stop_gradient is needed to exclude tf.maximum from derivative 
        step_func = sign + tf.stop_gradient(tf.maximum(0.0, sign) - sign) 
        return step_func 

在我的实现中有一个警告:tf.sign(0) 返回零值所以 heaviside(0) 也返回零，我想要 heaviside( 0) 返回 1。我怎样才能实现这种行为？

请您参考如下方法:

一个非常hacky的方法是使用

1 - max(0.0, sign(-x))  

作为你的步骤函数而不是

max(0.0, sign(x)) 

另一种选择是使用 greater_equal 并将结果转换为您想要的类型，并使用您已有的 sigmoid 覆盖覆盖其梯度。