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

Использование аннотаций типов для повышения производительности

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

@tf.function(experimental_follow_type_hints=True)
def f_with_hints(x: tf.Tensor):
  print('Tracing')
  return x
@tf.function(experimental_follow_type_hints=False)
def f_no_hints(x: tf.Tensor):
  print('Tracing')
  return x
f_no_hints(1)
Tracing
<tf.Tensor: shape=(), dtype=int32, numpy=1>
f_no_hints(2)
Tracing
<tf.Tensor: shape=(), dtype=int32, numpy=2>
f_with_hints(1)
Tracing
<tf.Tensor: shape=(), dtype=int32, numpy=1>
f_with_hints(2)
<tf.Tensor: shape=(), dtype=int32, numpy=2>