在TensorFlow中实现heaviside step函数答案

【问题标题】：Implementing heaviside step function in TensorFlow在TensorFlow中实现heaviside step函数
【发布时间】：2018-09-22 15:07:02
【问题描述】：

我想在 TensorFlow 中创建重载阶梯函数。由于 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。我怎样才能实现这种行为？

【问题讨论】：

你知道tensorflow.org/api_docs/python/tf/greater_equal吗？
是的，我考虑过使用greater_equal。问题是greater_equal 返回 bool Tensor，我不确定如何在向前和向后的步骤中处理浮点到布尔转换。
对于前进，您只需施放，对于后退，您将使用已有的内容覆盖它。

标签： python tensorflow machine-learning

【解决方案1】：

一个非常hacky的方法是使用

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

作为你的阶梯函数而不是

max(0.0, sign(x))

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

【讨论】：

【解决方案2】：

对您的代码最简单的修复方法是在tf.sign() 的结果中添加一个小数字，然后再次获取符号。这将导致 0 得到 1：

sign = tf.sign ( tf.sign( x ) + 0.1 )

【讨论】：

【解决方案3】：

好吧，我想我明白了。非常感谢etarion 指出了解决我问题的正确方法。

所以基本思想是使用tf.greater_equal，而不是tf.sign和maximum的组合。自定义渐变应用于tf.identity 操作。

heaviside 函数的更新实现如下：

import tensorflow as tf

@tf.RegisterGradient("HeavisideGrad")
def _heaviside_grad(unused_op: tf.Operation, grad: tf.Tensor):
    return tf.maximum(0.0, 1.0 - tf.abs(unused_op.inputs[0])) * grad


def heaviside(x: tf.Tensor, g: tf.Graph = tf.get_default_graph()):
    custom_grads = {
        "Identity": "HeavisideGrad"
    }
    with g.gradient_override_map(custom_grads):
        i = tf.identity(x, name="identity_" + str(uuid.uuid1()))
        ge = tf.greater_equal(x, 0, name="ge_" + str(uuid.uuid1()))
        # tf.stop_gradient is needed to exclude tf.to_float from derivative
        step_func = i + tf.stop_gradient(tf.to_float(ge) - i)
        return step_func

【讨论】：

【解决方案4】：

这将使单位步长函数，仅使用 TensorFlow API，因此结果仍然是张量：

#in Eager mode
def heaviside(v):
  return 1-tf.reduce_max(tf.constant([0,-tf.sign(v).numpy()], tf.float32));

在 TensorFlow 2 中，最好使用装饰器 @tf.custom_gradient：

@tf.custom_gradient
def heaviside(X):
  #This custom op is converted to graph, no 'if', 'else' allowed,
  #so use 'tf.cond'
  List = [];

  for I in range(BSIZE): #Batch size
    Item = tf.cond(X[I]<0, lambda: tf.constant([0], tf.float32), 
                           lambda: tf.constant([1], tf.float32));  
    List.append(Item);

  U = tf.stack(List);

  #Heaviside half-maximum formula
  #U = (tf.sign(X)+1)/2;

  #Div is differentiation intermediate value
  def grad(Div):
    return Div*1; #Heaviside has no gradient, use 1.

  return U,grad;

【讨论】：