python - 如何在 tensorflow 中的张量上自定义逐元素函数？

Question

假设我有一个张量，它可能包含正值和负值:

[ 1, -1, 2, -2 ]

现在，我想对正值应用 log(x)，对负值应用常量 -10:

[ log(1), -10, log(2), -10 ]

换句话说，我想要一个像numpy.vectorize这样的函数。这在 tensorflow 中可能吗？

一种可能的方法是使用不可学习的变量，但我不知道它是否可以正确地进行反向传播。

Answer 1

tf.map_fn()使您能够跨向量元素(或高维张量的切片)映射任意 TensorFlow 子计算。例如:

a = tf.constant([1.0, -1.0, 2.0, -2.0])

def f(elem):
  return tf.where(elem > 0, tf.log(elem), -10.0)

  # Alternatively, if the computation is more expensive than `tf.log()`, use
  # `tf.cond()` to ensure that only one branch is executed:
  # return tf.where(elem > 0, lambda: tf.log(elem), lambda: -10.0)

result = tf.map_fn(f, a)