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。

• 結果函數被假定為有狀態的並且永遠不會被優化。