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 参数。

由于它只被跟踪一次，因此变量和状态可以在函数内部创建并由函数包装对象拥有。