tf.function

Компилирует функцию в вызываемый граф TensorFlow. (устаревшие аргументы) (устаревшие аргументы) (устаревшие аргументы)

tf.function(
    func=None,
    input_signature=None,
    autograph=True,
    jit_compile=None,
    reduce_retracing=False,
    experimental_implements=None,
    experimental_autograph_options=None,
    experimental_attributes=None,
    experimental_relax_shapes=None,
    experimental_compile=None,
    experimental_follow_type_hints=None
) -> tf.types.experimental.PolymorphicFunction

Используется в ноутбуках

tf.function создает tf.types.experimental.PolymorphicFunction, который выполняет граф TensorFlow (tf.Graph), созданный путем трассировочной компиляции операций TensorFlow в func. Дополнительную информацию об этом можно найти в Введении в графы и tf.function.

См. Улучшенную производительность с tf.function для получения советов по производительности и известных ограничений.

Пример использования:

@tf.function
def f(x, y):
  return x ** 2 + y
x = tf.constant([2, 3])
y = tf.constant([3, -2])
f(x, y)
<tf.Tensor: ... numpy=array([7, 7], ...)>

Трассировочная компиляция позволяет выполнять операции, не являющиеся операциями TensorFlow, но при особых условиях. В целом, только операции TensorFlow гарантированно выполняются и создают новые результаты всякий раз, когда вызывается PolymorphicFunction.

Функции

func может использовать операторы Python с зависимостью от данных, включая if, for, while break, continue и return:

@tf.function
def f(x):
  if tf.reduce_sum(x) > 0:
    return x * x
  else:
    return -x // 2
f(tf.constant(-2))
<tf.Tensor: ... numpy=1>

Закрытие func может включать объекты tf.Tensor и tf.Variable:

@tf.function
def f():
  return x ** 2 + y
x = tf.constant([-2, -3])
y = tf.Variable([3, -2])
f()
<tf.Tensor: ... numpy=array([7, 7], ...)>

func также может использовать операции с побочными эффектами, такие как tf.print, tf.Variable и другие:

v = tf.Variable(1)
@tf.function
def f(x):
  for i in tf.range(x):
    v.assign_add(i)
f(3)
v
<tf.Variable ... numpy=4>

l = []
@tf.function
def f(x):
  for i in x:
    l.append(i + 1)    # Caution! Will only happen once when tracing
f(tf.constant([1, 2, 3]))
l
[<tf.Tensor ...>]

Вместо этого используйте коллекции TensorFlow, такие как tf.TensorArray:

@tf.function
def f(x):
  ta = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
  for i in range(len(x)):
    ta = ta.write(i, x[i] + 1)
  return ta.stack()
f(tf.constant([1, 2, 3]))
<tf.Tensor: ..., numpy=array([2, 3, 4], ...)>

tf.function создает полиморфные вызываемые функции

Внутренне, tf.types.experimental.PolymorphicFunction может содержать несколько tf.types.experimental.ConcreteFunction, каждая из которых специализирована для аргументов с различными типами данных или формами, поскольку TensorFlow может выполнять больше оптимизаций над графами с определенными формами, типами данных и значениями постоянных аргументов. tf.function обрабатывает любые чистые значения Python как неявные объекты (лучше всего представляемые как константы времени компиляции) и строит отдельный tf.Graph для каждого набора аргументов Python, с которым он сталкивается. Более подробную информацию см. в руководстве tf.function.

Выполнение PolymorphicFunction выберет и выполнит соответствующую ConcreteFunction на основе типов и значений аргументов.

Чтобы получить отдельную ConcreteFunction, используйте метод PolymorphicFunction.get_concrete_function. Его можно вызывать с теми же аргументами, что и func, и он возвращает tf.types.experimental.ConcreteFunction. ConcreteFunction основаны на единственном tf.Graph:

@tf.function
def f(x):
  return x + 1
isinstance(f.get_concrete_function(1).graph, tf.Graph)
True

ConcreteFunction можно выполнять так же, как и PolymorphicFunction, но их вход ограничен типами, к которым они специализированы.

Переотслеживание

ConcreteFunctions строятся (отслеживаются) на лету, поскольку PolymorphicFunction вызывается с новыми типами или формами TensorFlow, или с новыми значениями Python в качестве аргументов. Когда PolymorphicFunction строит новую трассу, говорят, что func переотслеживается. Переотслеживание является частой проблемой производительности для tf.function, так как оно может быть значительно медленнее, чем выполнение уже отслеженного графа. Желательно минимизировать количество переотслеживаний в вашем коде.

@tf.function
def f(x):
  return tf.abs(x)
f1 = f.get_concrete_function(1)
f2 = f.get_concrete_function(2)  # Slow - compiles new graph
f1 is f2
False
f1 = f.get_concrete_function(tf.constant(1))
f2 = f.get_concrete_function(tf.constant(2))  # Fast - reuses f1
f1 is f2
True

Числовые аргументы Python следует использовать только тогда, когда они принимают несколько различных значений, таких как гиперпараметры, например, количество слоев в нейронной сети.

Подписи ввода

Для аргументов тензоров PolymorphicFunction создает новую ConcreteFunction для каждого уникального набора форм и типов входных данных. Приведенный ниже пример создает две отдельные ConcreteFunction, каждая специализированная для разной формы:

@tf.function
def f(x):
  return x + 1
vector = tf.constant([1.0, 1.0])
matrix = tf.constant([[3.0]])
f.get_concrete_function(vector) is f.get_concrete_function(matrix)
False

«Подпись ввода» может быть необязательно предоставлена tf.function для управления этим процессом. Подпись ввода определяет форму и тип каждого аргумента тензора в функцию с помощью объекта tf.TensorSpec. Можно использовать более общие формы. Это гарантирует создание только одной ConcreteFunction и ограничивает PolymorphicFunction заданными формами и типами. Это эффективный способ ограничить переотслеживание, когда тензоры имеют динамические формы.

@tf.function(
    input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
def f(x):
  return x + 1
vector = tf.constant([1.0, 1.0])
matrix = tf.constant([[3.0]])
f.get_concrete_function(vector) is f.get_concrete_function(matrix)
True

Переменные могут создаваться только один раз

tf.function разрешает создание новых объектов tf.Variable только при первом вызове:

class MyModule(tf.Module):
  def __init__(self):
    self.v = None

  @tf.function
  def __call__(self, x):
    if self.v is None:
      self.v = tf.Variable(tf.ones_like(x))
    return self.v * x

В общем случае рекомендуется создавать tf.Variable вне tf.function. В простых случаях сохранение состояния через границы tf.function можно реализовать с использованием чисто функционального стиля, в котором состояние представлено tf.Tensor , передаваемыми в качестве аргументов и возвращаемыми в качестве значений возврата.

Сравните два стиля ниже:

state = tf.Variable(1)
@tf.function
def f(x):
  state.assign_add(x)
f(tf.constant(2))  # Non-pure functional style
state
<tf.Variable ... numpy=3>

state = tf.constant(1)
@tf.function
def f(state, x):
  state += x
  return state
state = f(state, tf.constant(2))  # Pure functional style
state
<tf.Tensor: ... numpy=3>

Операции Python выполняются только один раз на отслеживание

func может содержать операции TensorFlow, смешанные с чистыми операциями Python. Однако при выполнении функции будут выполняться только операции TensorFlow. Операции Python выполняются только один раз во время отслеживания. Если операции TensorFlow зависят от результатов операций Python, эти результаты будут заморожены в графе.

@tf.function
def f(a, b):
  print('this runs at trace time; a is', a, 'and b is', b)
  return b
f(1, tf.constant(1))
this runs at trace time; a is 1 and b is Tensor("...", shape=(), dtype=int32)
<tf.Tensor: shape=(), dtype=int32, numpy=1>

f(1, tf.constant(2))
<tf.Tensor: shape=(), dtype=int32, numpy=2>

f(2, tf.constant(1))
this runs at trace time; a is 2 and b is Tensor("...", shape=(), dtype=int32)
<tf.Tensor: shape=(), dtype=int32, numpy=1>

f(2, tf.constant(2))
<tf.Tensor: shape=(), dtype=int32, numpy=2>