Python tf.compat.v1.wrap_function用法及代碼示例

將 TF 1.x 函數 fn 包裝成圖形函數。

用法

tf.compat.v1.wrap_function(
    fn, signature, name=None
)

python 函數 fn 將使用 signature 中指定的符號參數調用一次，跟蹤並轉換為圖形函數。 fn 創建的任何變量都將歸 wrap_function 返回的對象所有。可以使用與簽名匹配的張量調用生成的圖函數。

def f(x, do_add):
  v = tf.Variable(5.0)
  if do_add:
    op = v.assign_add(x)
  else:
    op = v.assign_sub(x)
  with tf.control_dependencies([op]):
    return v.read_value()

f_add = tf.compat.v1.wrap_function(f, [tf.TensorSpec((), tf.float32), True])

assert float(f_add(1.0)) == 6.0
assert float(f_add(1.0)) == 7.0

# Can call tf.compat.v1.wrap_function again to get a new trace, a new set
# of variables, and possibly different non-template arguments.
f_sub= tf.compat.v1.wrap_function(f, [tf.TensorSpec((), tf.float32), False])

assert float(f_sub(1.0)) == 4.0
assert float(f_sub(1.0)) == 3.0

tf.compat.v1.wrap_function 和 tf.function 都創建了一個可調用的 TensorFlow 圖。但是，雖然 tf.function 運行所有有狀態操作(例如 tf.print )和序列操作以提供與即刻執行相同的語義，但 wrap_function 更接近於 TensorFlow 1.x 中 session.run 的行為。它不會運行任何操作，除非它們需要通過數據依賴或控製依賴來計算函數的輸出。它也不會對操作進行排序。

與 tf.function 不同，wrap_function 隻會跟蹤 Python 函數一次。與 TF 1.x 中的占位符一樣，形狀和數據類型必須提供給 wrap_function 的 signature 參數。

由於它隻被跟蹤一次，因此變量和狀態可以在函數內部創建並由函數包裝對象擁有。