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 ближе к поведению session.run в TensorFlow 1.x. Она не будет выполнять никаких операций, если они не нужны для вычисления выходных данных функции, либо через зависимость данных, либо через зависимость управления. Она также не будет последовательно выполнять операции.

В отличие от tf.function, wrap_function будет отслеживать функцию Python только один раз. Как и с плейсхолдерами в TF 1.x, формы и типы данных должны быть указаны в аргументе signature функции wrap_function.

Поскольку отслеживание выполняется только один раз, переменные и состояние могут быть созданы внутри функции и принадлежать объекту-обертке функции.