tf.function

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

tf.function(
    func=None, input_signature=None, autograph=True, experimental_implements=None,
    experimental_autograph_options=None, experimental_relax_shapes=False,
    experimental_compile=None
)

tf.function создаёт вызываемый объект, который выполняет граф TensorFlow (tf.Graph), созданный путем трассировки компиляции операций TensorFlow в func, эффективно выполняя func как граф TensorFlow.

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

@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], ...)>

Особенности

func может использовать условную логику, зависящую от данных, включая 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.function может создавать более одного графа для поддержки аргументов с различными типами и формами данных, так как TensorFlow может создавать более эффективные графы, специализированные на формах и типах данных. tf.function также обрабатывает любые чистые значения Python как непрозрачные объекты и строит отдельный граф для каждого набора аргументов Python, которые он встречает.

Чтобы получить отдельный граф, используйте метод get_concrete_function вызываемого объекта, созданного tf.function. Он может быть вызван с теми же аргументами, что и func, и возвращает специальный объект tf.Graph:

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

@tf.function
def f(x):
  return tf.abs(x)
f1 = f.get_concrete_function(1)
f2 = f.get_concrete_function(2)  # Slow - builds 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 следует использовать только тогда, когда они принимают несколько различных значений, таких как гиперпараметры, например, количество слоёв в нейронной сети.

Подписи входных данных

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

@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. Можно использовать более общие формы. Это полезно для предотвращения создания нескольких графов, когда тензоры имеют динамические формы. Это также ограничивает форму и тип данных тензоров, которые могут быть использованы:

@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 и передавать их в качестве аргументов.