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 следует использовать только в том случае, если они принимают небольшое число различных значений, например гиперпараметры, такие как количество слоев нейронной сети.

Сигнатуры входных данных

Для аргументов типа Tensor 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 можно дополнительно передать «сигнатуру входных данных». Сигнатура входных данных задает форму и тип каждого аргумента Tensor функции с помощью объекта 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>