Python tf.numpy_function用法及代码示例

包装一个 python 函数并将其用作 TensorFlow 操作。

用法

tf.numpy_function(
    func, inp, Tout, name=None
)

给定一个 python 函数 func 将此函数包装为 TensorFlow 函数中的一个操作。 func 必须将 numpy 数组作为其参数并返回 numpy 数组作为其输出。

以下示例使用 np.sinh() 作为图中的操作创建 TensorFlow 图：

def my_numpy_func(x):
  # x will be a numpy array with the contents of the input to the
  # tf.function
  return np.sinh(x)
@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
def tf_function(input):
  y = tf.numpy_function(my_numpy_func, [input], tf.float32)
  return y * y
tf_function(tf.constant(1.))
<tf.Tensor:shape=(), dtype=float32, numpy=1.3810978>

与 tf.py_function 的比较：tf.py_function 和 tf.numpy_function 非常相似，除了 tf.numpy_function 采用 numpy 数组，而不是 tf.Tensor s。如果您希望函数包含 tf.Tensors ，并且在函数中执行的任何 TensorFlow 操作都是可微的，请使用 tf.py_function 。

注意：tf.numpy_function 操作具有以下已知限制：

• 函数体(即 func )不会在 tf.SavedModel 中序列化。因此，如果您需要序列化模型并在不同的环境中恢复它，则不应使用此函数。

• 该操作必须在与调用 tf.numpy_function() 的 Python 程序相同的地址空间中运行。如果您使用分布式 TensorFlow，则必须在与调用 tf.numpy_function 的程序相同的进程中运行 tf.distribute.Server，您必须将创建的操作固定到该服务器中的设备(例如，使用 with tf.device(): )。

• 由于该函数采用 numpy 数组，因此您不能通过 numpy_function 获取渐变。如果您需要可区分的东西，请考虑使用 tf.py_function。

• 结果函数被假定为有状态的并且永远不会被优化。