Python tf.py_function用法及代碼示例

將一個 python 函數包裝到一個 TensorFlow 操作中，該操作會即刻地執行它。

用法

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

此函數允許將 TensorFlow 圖中的計算表示為 Python 函數。特別是，它將 Python 函數 func 包裝在一次可微分的 TensorFlow 操作中，該操作在啟用即刻執行的情況下執行它。因此，tf.py_function 可以使用 Python 構造(if , while , for 等)而不是 TensorFlow 控製流構造(tf.cond、tf.while_loop)來表達控製流。例如，您可以使用tf.py_function 來實現日誌集線器函數：

def log_huber(x, m):
  if tf.abs(x) <= m:
    return x**2
  else:
    return m**2 * (1 - 2 * tf.math.log(m) + tf.math.log(x**2))

x = tf.compat.v1.placeholder(tf.float32)
m = tf.compat.v1.placeholder(tf.float32)

y = tf.py_function(func=log_huber, inp=[x, m], Tout=tf.float32)
dy_dx = tf.gradients(y, x)[0]

with tf.compat.v1.Session() as sess:
  # The session executes `log_huber` eagerly. Given the feed values below,
  # it will take the first branch, so `y` evaluates to 1.0 and
  # `dy_dx` evaluates to 2.0.
  y, dy_dx = sess.run([y, dy_dx], feed_dict={x:1.0, m:2.0})

您還可以使用tf.py_function 在運行時使用 Python 工具調試模型，即，您可以隔離要調試的代碼部分，將它們包裝在 Python 函數中並根據需要插入 pdb 跟蹤點或打印語句，以及將這些函數包裝在 tf.py_function 中。

有關 Eager Execution 的更多信息，請參閱 Eager 指南。

tf.py_function 在本質上與 tf.compat.v1.py_func 相似，但與後者不同的是，前者允許您在包裝的 Python 函數中使用 TensorFlow 操作。特別是，雖然tf.compat.v1.py_func 僅在 CPU 上運行並包裝以 NumPy 數組作為輸入並返回 NumPy 數組作為輸出的函數，但 tf.py_function 可以放置在 GPU 上並包裝以張量作為輸入的函數，在它們的主體中執行 TensorFlow 操作，並返回張量作為輸出。

與 tf.compat.v1.py_func 一樣，tf.py_function 在序列化和分發方掩碼有以下限製：

• 函數體(即 func )不會在 GraphDef 中序列化。因此，如果您需要序列化模型並在不同的環境中恢複它，則不應使用此函數。

• 該操作必須在與調用 tf.py_function() 的 Python 程序相同的地址空間中運行。如果您使用分布式 TensorFlow，則必須在與調用 tf.py_function() 的程序相同的進程中運行 tf.distribute.Server，並且必須將創建的操作固定到該服務器中的設備(例如，使用 with tf.device(): )。