Лучшая производительность с tf.function

Посмотреть на TensorFlow.org Запустить в Google Colab Посмотреть исходный код на GitHubСкачать блокнот

В TensorFlow 2 нетерпеливое выполнение включено по умолчанию. Пользовательский интерфейс интуитивно понятен и гибок (выполнение одноразовых операций намного проще и быстрее), но это может быть достигнуто за счет производительности и возможности развертывания.

Вы можете использовать tf.function для создания графиков из ваших программ. Это инструмент преобразования, который создает независимые от Python графы потоков данных из вашего кода Python. Это поможет вам создавать производительные и портативные модели, и это необходимо для использования SavedModel .

Это руководство поможет вам понять, как работает tf.function , чтобы вы могли эффективно его использовать.

Основные выводы и рекомендации:

  • Выполните отладку в активном режиме, а затем украсьте @tf.function .
  • Не полагайтесь на побочные эффекты Python, такие как мутация объекта или добавление списка.
  • tf.function лучше всего работает с операциями TensorFlow; Вызовы NumPy и Python преобразуются в константы.

Настраивать

import tensorflow as tf

Определите вспомогательную функцию, чтобы продемонстрировать типы ошибок, с которыми вы можете столкнуться:

import traceback
import contextlib

# Some helper code to demonstrate the kinds of errors you might encounter.
@contextlib.contextmanager
def assert_raises(error_class):
  try:
    yield
  except error_class as e:
    print('Caught expected exception \n  {}:'.format(error_class))
    traceback.print_exc(limit=2)
  except Exception as e:
    raise e
  else:
    raise Exception('Expected {} to be raised but no error was raised!'.format(
        error_class))

Основы

использование

Function , которую вы определяете (например, применяя декоратор @tf.function ), аналогична основной операции TensorFlow: вы можете выполнить ее с нетерпением; вы можете вычислять градиенты; и так далее.

@tf.function  # The decorator converts `add` into a `Function`.
def add(a, b):
  return a + b

add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]

<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[2., 2.],
       [2., 2.]], dtype=float32)>

v = tf.Variable(1.0)
with tf.GradientTape() as tape:
  result = add(v, 1.0)
tape.gradient(result, v)

<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

Вы можете использовать Function внутри других Function .

@tf.function
def dense_layer(x, w, b):
  return add(tf.matmul(x, w), b)

dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))

<tf.Tensor: shape=(3, 2), dtype=float32, numpy=
array([[3., 3.],
       [3., 3.],
       [3., 3.]], dtype=float32)>

Function могут быть быстрее, чем энергичный код, особенно для графов с большим количеством небольших операций. Но для графов с несколькими дорогостоящими операциями (например, свертки) вы можете не увидеть большого ускорения.

import timeit
conv_layer = tf.keras.layers.Conv2D(100, 3)

@tf.function
def conv_fn(image):
  return conv_layer(image)

image = tf.zeros([1, 200, 200, 100])
# Warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Note how there's not much difference in performance for convolutions")

Eager conv: 0.006058974999177735
Function conv: 0.005791576000774512
Note how there's not much difference in performance for convolutions

Отслеживание

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

Что такое "отслеживание"?

Function запускает вашу программу в TensorFlow Graph . Однако tf.Graph не может представлять все то, что вы написали бы в энергичной программе TensorFlow. Например, Python поддерживает полиморфизм, но tf.Graph требует, чтобы его входные данные имели указанный тип данных и размерность. Или вы можете выполнять побочные задачи, такие как чтение аргументов командной строки, создание ошибки или работа с более сложным объектом Python; ни одна из этих вещей не может работать в tf.Graph .

Function устраняет этот пробел, разделяя ваш код на два этапа:

1) На первом этапе, называемом « трассировка », Function создает новый tf.Graph . Код Python работает нормально, но все операции TensorFlow (например, добавление двух тензоров) откладываются : они захватываются tf.Graph и не выполняются.

2) На втором этапе tf.Graph , содержащий все, что было отложено на первом этапе. Этот этап намного быстрее, чем этап трассировки.

В зависимости от своих входных данных Function не всегда будет запускать первую стадию при вызове. См. «Правила отслеживания» ниже, чтобы лучше понять, как он определяет это. Пропуск первого этапа и выполнение только второго этапа — вот что обеспечивает высокую производительность TensorFlow.

Когда Function принимает решение о трассировке, за этапом трассировки сразу же следует второй этап, поэтому вызов Function одновременно создает и запускает tf.Graph . Позже вы увидите, как с помощью get_concrete_function можно запустить только этап трассировки.

Когда вы передаете аргументы разных типов в Function , выполняются оба этапа:

@tf.function
def double(a):
  print("Tracing with", a)
  return a + a

print(double(tf.constant(1)))
print()
print(double(tf.constant(1.1)))
print()
print(double(tf.constant("a")))
print()

Tracing with Tensor("a:0", shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Tracing with Tensor("a:0", shape=(), dtype=float32)
tf.Tensor(2.2, shape=(), dtype=float32)

Tracing with Tensor("a:0", shape=(), dtype=string)
tf.Tensor(b'aa', shape=(), dtype=string)

Обратите внимание, что если вы неоднократно вызываете Function с одним и тем же типом аргумента, TensorFlow пропустит этап трассировки и повторно использует ранее отслеженный график, поскольку сгенерированный график будет идентичным.

# This doesn't print 'Tracing with ...'
print(double(tf.constant("b")))

tf.Tensor(b'bb', shape=(), dtype=string)

Вы можете использовать pretty_printed_concrete_signatures() , чтобы увидеть все доступные трассировки:

print(double.pretty_printed_concrete_signatures())

double(a)
  Args:
    a: int32 Tensor, shape=()
  Returns:
    int32 Tensor, shape=()

double(a)
  Args:
    a: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()

double(a)
  Args:
    a: string Tensor, shape=()
  Returns:
    string Tensor, shape=()

До сих пор вы видели, что tf.function создает кешированный динамический уровень диспетчеризации поверх логики трассировки графа TensorFlow. Чтобы уточнить терминологию:

  • tf.Graph — это необработанное, независимое от языка, переносимое представление вычислений TensorFlow.
  • ConcreteFunction обертывает tf.Graph .
  • Function управляет кешем ConcreteFunction и выбирает правильную для ваших входных данных.
  • tf.function оборачивает функцию Python, возвращая объект Function .
  • Tracing создает tf.Graph и заключает его в ConcreteFunction , также известную как трассировка.

Правила розыска

Function определяет, следует ли повторно использовать трассируемую ConcreteFunction , вычисляя ключ кэша из входных аргументов и kwargs. Ключ кэша — это ключ, который идентифицирует ConcreteFunction на основе входных аргументов и kwargs вызова Function в соответствии со следующими правилами (которые могут измениться):

  • Ключ, сгенерированный для tf.Tensor , — это его форма и тип.
  • Ключ, сгенерированный для tf.Variable , представляет собой уникальный идентификатор переменной.
  • Ключ, сгенерированный для примитива Python (например, int , float , str ), является его значением.
  • Ключ, сгенерированный для вложенных dict s, list s, tuple s, namedtuple s и attr s, представляет собой сглаженный кортеж leaf-keys (см. nest.flatten ). (В результате этого выравнивания вызов конкретной функции со структурой вложенности, отличной от той, которая использовалась во время трассировки, приведет к ошибке TypeError).
  • Для всех других типов Python ключ уникален для объекта. Таким образом, функция или метод отслеживаются независимо для каждого экземпляра, с которым они вызываются.

Управление откатом

Повторная трассировка, когда ваша Function создает более одной трассировки, помогает гарантировать, что TensorFlow генерирует правильные графики для каждого набора входных данных. Однако трассировка — дорогостоящая операция! Если ваша Function отслеживает новый график для каждого вызова, вы обнаружите, что ваш код выполняется медленнее, чем если бы вы не использовали tf.function .

Для управления поведением трассировки можно использовать следующие методы:

  • Укажите input_signature в tf.function , чтобы ограничить трассировку.
@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)

print(next_collatz(tf.constant([1, 2])))
# You specified a 1-D tensor in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([[1, 2], [3, 4]]))

# You specified an int32 dtype in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([1.0, 2.0]))

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([4 1], shape=(2,), dtype=int32)
Caught expected exception 
  <class 'ValueError'>:
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/1851403433.py", line 9, in <module>
    next_collatz(tf.constant([[1, 2], [3, 4]]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None)).
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/1851403433.py", line 13, in <module>
    next_collatz(tf.constant([1.0, 2.0]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor([1. 2.], shape=(2,), dtype=float32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None)).

  • Укажите размерность [None] в tf.TensorSpec , чтобы обеспечить гибкость повторного использования трассировки.

    Поскольку TensorFlow сопоставляет тензоры на основе их формы, использование размера None в качестве подстановочного знака позволит Function повторно использовать трассировки для входных данных переменного размера. Ввод переменного размера может происходить, если у вас есть последовательности разной длины или изображения разных размеров для каждого пакета (см., например, учебные пособия Transformer и Deep Dream ).

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def g(x):
  print('Tracing with', x)
  return x

# No retrace!
print(g(tf.constant([1, 2, 3])))
print(g(tf.constant([1, 2, 3, 4, 5])))

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([1 2 3], shape=(3,), dtype=int32)
tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)

  • Приведите аргументы Python к тензорам, чтобы уменьшить повторную трассировку.

    Часто аргументы Python используются для управления гиперпараметрами и построениями графиков — например, num_layers=10 или training=True или nonlinearity='relu' . Таким образом, если аргумент Python изменится, имеет смысл пересмотреть график.

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

def train_one_step():
  pass

@tf.function
def train(num_steps):
  print("Tracing with num_steps = ", num_steps)
  tf.print("Executing with num_steps = ", num_steps)
  for _ in tf.range(num_steps):
    train_one_step()

print("Retracing occurs for different Python arguments.")
train(num_steps=10)
train(num_steps=20)

print()
print("Traces are reused for Tensor arguments.")
train(num_steps=tf.constant(10))
train(num_steps=tf.constant(20))

Retracing occurs for different Python arguments.
Tracing with num_steps =  10
Executing with num_steps =  10
Tracing with num_steps =  20
Executing with num_steps =  20

Traces are reused for Tensor arguments.
Tracing with num_steps =  Tensor("num_steps:0", shape=(), dtype=int32)
Executing with num_steps =  10
Executing with num_steps =  20

Если вам нужно принудительно вернуться, создайте новую Function . Гарантируется, что отдельные объекты Function не будут совместно использовать трассировки.

def f():
  print('Tracing!')
  tf.print('Executing')

tf.function(f)()
tf.function(f)()

Tracing!
Executing
Tracing!
Executing

Получение конкретных функций

Каждый раз при трассировке функции создается новая конкретная функция. Вы можете напрямую получить конкретную функцию, используя get_concrete_function .

print("Obtaining concrete trace")
double_strings = double.get_concrete_function(tf.constant("a"))
print("Executing traced function")
print(double_strings(tf.constant("a")))
print(double_strings(a=tf.constant("b")))

Obtaining concrete trace
Executing traced function
tf.Tensor(b'aa', shape=(), dtype=string)
tf.Tensor(b'bb', shape=(), dtype=string)

# You can also call get_concrete_function on an InputSpec
double_strings_from_inputspec = double.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.string))
print(double_strings_from_inputspec(tf.constant("c")))

tf.Tensor(b'cc', shape=(), dtype=string)

При печати ConcreteFunction отображается сводка ее входных аргументов (с типами) и тип вывода.

print(double_strings)

ConcreteFunction double(a)
  Args:
    a: string Tensor, shape=()
  Returns:
    string Tensor, shape=()

Вы также можете напрямую получить сигнатуру конкретной функции.

print(double_strings.structured_input_signature)
print(double_strings.structured_outputs)

((TensorSpec(shape=(), dtype=tf.string, name='a'),), {})
Tensor("Identity:0", shape=(), dtype=string)

Использование конкретной трассировки с несовместимыми типами вызовет ошибку

with assert_raises(tf.errors.InvalidArgumentError):
  double_strings(tf.constant(1))

Caught expected exception 
  <class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/3196284684.py", line 2, in <module>
    double_strings(tf.constant(1))
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute __inference_double_162 as input #0(zero-based) was expected to be a string tensor but is a int32 tensor [Op:__inference_double_162]

Вы можете заметить, что аргументы Python имеют специальную обработку в входной сигнатуре конкретной функции. До TensorFlow 2.3 аргументы Python просто удалялись из сигнатуры конкретной функции. Начиная с TensorFlow 2.3, аргументы Python остаются в подписи, но они ограничены тем, что принимают значение, установленное во время трассировки.

@tf.function
def pow(a, b):
  return a ** b

square = pow.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)
print(square)

ConcreteFunction pow(a, b=2)
  Args:
    a: float32 Tensor, shape=<unknown>
  Returns:
    float32 Tensor, shape=<unknown>

assert square(tf.constant(10.0)) == 100

with assert_raises(TypeError):
  square(tf.constant(10.0), b=3)

Caught expected exception 
  <class 'TypeError'>:
Traceback (most recent call last):
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1721, in _call_impl
    cancellation_manager)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1765, in _call_with_flat_signature
    raise TypeError(f"{self._flat_signature_summary()} got unexpected "
TypeError: pow(a) got unexpected keyword arguments: b.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/2310937119.py", line 4, in <module>
    square(tf.constant(10.0), b=3)
TypeError: ConcreteFunction pow(a, b) was constructed with int value 2 in b, but was called with int value 3.

Получение графиков

Каждая конкретная функция является вызываемой оболочкой вокруг tf.Graph . Хотя извлечение фактического объекта tf.Graph обычно не требуется, вы можете легко получить его из любой конкретной функции.

graph = double_strings.graph
for node in graph.as_graph_def().node:
  print(f'{node.input} -> {node.name}')

[] -> a
['a', 'a'] -> add
['add'] -> Identity

Отладка

В общем случае отладка кода в активном режиме проще, чем внутри tf.function . Вы должны убедиться, что ваш код выполняется без ошибок в активном режиме, прежде чем украшать его с помощью tf.function . Чтобы помочь в процессе отладки, вы можете вызвать tf.config.run_functions_eagerly(True) для глобального отключения и повторного включения tf.function .

При отслеживании проблем, которые появляются только в tf.function , вот несколько советов:

  • Обычные старые вызовы print Python выполняются только во время трассировки, помогая вам отслеживать, когда ваша функция (повторно) трассируется.
  • tf.print будут выполняться каждый раз и могут помочь вам отслеживать промежуточные значения во время выполнения.
  • tf.debugging.enable_check_numerics — это простой способ отследить, где создаются NaN и Inf.
  • pdb ( отладчик Python ) может помочь вам понять, что происходит во время трассировки. (Предупреждение: pdb перебросит вас в исходный код, преобразованный AutoGraph.)

Преобразования автографа

AutoGraph — это библиотека, которая включена по умолчанию в tf.function и преобразует подмножество нетерпеливого кода Python в граф-совместимые операции TensorFlow. Это включает в себя поток управления, например if , for , while .

Операции TensorFlow, такие как tf.cond и tf.while_loop , продолжают работать, но поток управления часто легче написать и понять, когда он написан на Python.

# A simple loop

@tf.function
def f(x):
  while tf.reduce_sum(x) > 1:
    tf.print(x)
    x = tf.tanh(x)
  return x

f(tf.random.uniform([5]))

[0.666458249 0.713946581 0.723879576 0.330758929 0.184087753]
[0.582645297 0.613145649 0.619306684 0.319202513 0.182036072]
[0.524585426 0.546337605 0.550645113 0.308785647 0.18005164]
[0.481231302 0.497770309 0.501003504 0.299331933 0.178130865]
[0.447229207 0.460361809 0.462906033 0.290701121 0.176270396]
[0.419618756 0.430379033 0.432449728 0.282779962 0.174467146]
[0.396609187 0.405638 0.407366514 0.275476 0.172718227]
[0.377043903 0.384762734 0.386234313 0.268712848 0.17102097]
[0.360137492 0.366836458 0.368109286 0.262426734 0.169372901]
[0.345335096 0.351221472 0.352336824 0.256563932 0.167771652]
[0.332231969 0.337458342 0.338446289 0.251078814 0.166215062]
[0.320524871 0.325206399 0.326089561 0.24593246 0.164701089]
[0.309981436 0.314206958 0.31500268 0.241091311 0.163227797]
[0.300420195 0.304259449 0.304981351 0.236526251 0.161793426]
[0.291697085 0.295205742 0.295864582 0.232211992 0.160396278]
[0.283696055 0.286919087 0.287523568 0.228126258 0.159034774]
[0.276322395 0.279296666 0.27985391 0.224249557 0.157707423]
[0.269497961 0.272254 0.272769839 0.220564634 0.15641281]
[0.263157606 0.265720904 0.266200244 0.21705614 0.155149609]
[0.257246554 0.259638608 0.260085613 0.213710397 0.153916568]
[0.251718313 0.25395745 0.254375577 0.210515186 0.152712509]
[0.246533215 0.248635098 0.249027327 0.207459539 0.151536316]
[0.241657034 0.243635193 0.244004101 0.204533577 0.15038693]
[0.237060249 0.238926381 0.239274174 0.201728329 0.149263337]
[0.232717097 0.234481394 0.234810054 0.199035719 0.148164615]
[0.228605017 0.230276451 0.230587661 0.196448416 0.147089839]
[0.224704206 0.226290658 0.22658591 0.193959698 0.14603813]
[0.220997125 0.222505584 0.222786173 0.191563457 0.145008713]
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.21746822, 0.21890487, 0.21917202, 0.18925412, 0.14400077],
      dtype=float32)>

Если вам интересно, вы можете проверить код, сгенерированный автографом.

print(tf.autograph.to_code(f.python_function))

def tf__f(x):
    with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()

        def get_state():
            return (x,)

        def set_state(vars_):
            nonlocal x
            (x,) = vars_

        def loop_body():
            nonlocal x
            ag__.converted_call(ag__.ld(tf).print, (ag__.ld(x),), None, fscope)
            x = ag__.converted_call(ag__.ld(tf).tanh, (ag__.ld(x),), None, fscope)

        def loop_test():
            return (ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) > 1)
        ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('x',), {})
        try:
            do_return = True
            retval_ = ag__.ld(x)
        except:
            do_return = False
            raise
        return fscope.ret(retval_, do_return)

Условные

AutoGraph преобразует некоторые операторы if <condition> в эквивалентные вызовы tf.cond . Эта замена производится, если <condition> является тензором. В противном случае оператор if выполняется как условное выражение Python.

Условное выражение Python выполняется во время трассировки, поэтому к графу будет добавлена ​​ровно одна ветвь условного выражения. Без AutoGraph этот трассируемый граф не смог бы перейти на альтернативную ветвь, если существует поток управления, зависящий от данных.

tf.cond отслеживает и добавляет в граф обе ветви условного оператора, динамически выбирая ветвь во время выполнения. Трассировка может иметь непреднамеренные побочные эффекты; ознакомьтесь с эффектами трассировки AutoGraph для получения дополнительной информации.

@tf.function
def fizzbuzz(n):
  for i in tf.range(1, n + 1):
    print('Tracing for loop')
    if i % 15 == 0:
      print('Tracing fizzbuzz branch')
      tf.print('fizzbuzz')
    elif i % 3 == 0:
      print('Tracing fizz branch')
      tf.print('fizz')
    elif i % 5 == 0:
      print('Tracing buzz branch')
      tf.print('buzz')
    else:
      print('Tracing default branch')
      tf.print(i)

fizzbuzz(tf.constant(5))
fizzbuzz(tf.constant(20))

Tracing for loop
Tracing fizzbuzz branch
Tracing fizz branch
Tracing buzz branch
Tracing default branch
1
2
fizz
4
buzz
1
2
fizz
4
buzz
fizz
7
8
fizz
buzz
11
fizz
13
14
fizzbuzz
16
17
fizz
19
buzz

Дополнительные ограничения на операторы if, преобразованные с помощью AutoGraph, см. в справочной документации .

Петли

AutoGraph преобразует некоторые операторы for и while в эквивалентные циклические операции TensorFlow, такие как tf.while_loop . Если не преобразовать, цикл for или while выполняется как цикл Python.

Такая замена производится в следующих случаях:

  • for x in y : если y является тензором, преобразуйте в tf.while_loop . В особом случае, когда y является tf.data.Dataset , генерируется комбинация операций tf.data.Dataset .
  • while <condition> : если <condition> является тензором, преобразуйте в tf.while_loop .

Цикл Python выполняется во время трассировки, добавляя дополнительные операции в tf.Graph для каждой итерации цикла.

Цикл TensorFlow отслеживает тело цикла и динамически выбирает, сколько итераций выполнять во время выполнения. Тело цикла появляется только один раз в сгенерированном tf.Graph .

Дополнительные ограничения на операторы for и while , преобразованные с помощью AutoGraph, см. в справочной документации .

Перебор данных Python

Распространенной ошибкой является циклическое перебор данных Python/NumPy внутри tf.function . Этот цикл будет выполняться во время процесса трассировки, добавляя копию вашей модели в tf.Graph для каждой итерации цикла.

Если вы хотите обернуть весь цикл обучения в tf.function , самый безопасный способ сделать это — обернуть ваши данные как tf.data.Dataset , чтобы AutoGraph динамически разворачивал цикл обучения.

def measure_graph_size(f, *args):
  g = f.get_concrete_function(*args).graph
  print("{}({}) contains {} nodes in its graph".format(
      f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node)))

@tf.function
def train(dataset):
  loss = tf.constant(0)
  for x, y in dataset:
    loss += tf.abs(y - x) # Some dummy computation.
  return loss

small_data = [(1, 1)] * 3
big_data = [(1, 1)] * 10
measure_graph_size(train, small_data)
measure_graph_size(train, big_data)

measure_graph_size(train, tf.data.Dataset.from_generator(
    lambda: small_data, (tf.int32, tf.int32)))
measure_graph_size(train, tf.data.Dataset.from_generator(
    lambda: big_data, (tf.int32, tf.int32)))

train([(1, 1), (1, 1), (1, 1)]) contains 11 nodes in its graph
train([(1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]) contains 32 nodes in its graph
train(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 6 nodes in its graph
train(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 6 nodes in its graph

При переносе данных Python/NumPy в набор данных помните о tf.data.Dataset.from_generator и tf.data.Dataset.from_tensors . Первый будет хранить данные в Python и извлекать их через tf.py_function , что может повлиять на производительность, тогда как второй будет объединять копию данных в виде одного большого tf.constant() в графе, что может иметь последствия для памяти.

Чтение данных из файлов через TFRecordDataset , CsvDataset и т. д. — наиболее эффективный способ потребления данных, так как тогда TensorFlow сам может управлять асинхронной загрузкой и предварительной выборкой данных без необходимости привлечения Python. Чтобы узнать больше, см. руководство tf.data : сборка входных конвейеров TensorFlow .

Накопление значений в цикле

Распространенным шаблоном является накопление промежуточных значений из цикла. Обычно это достигается добавлением к списку Python или добавлением записей в словарь Python. Однако, поскольку это побочные эффекты Python, они не будут работать должным образом в динамически развернутом цикле. Используйте tf.TensorArray для накопления результатов динамически развернутого цикла.

batch_size = 2
seq_len = 3
feature_size = 4

def rnn_step(inp, state):
  return inp + state

@tf.function
def dynamic_rnn(rnn_step, input_data, initial_state):
  # [batch, time, features] -> [time, batch, features]
  input_data = tf.transpose(input_data, [1, 0, 2])
  max_seq_len = input_data.shape[0]

  states = tf.TensorArray(tf.float32, size=max_seq_len)
  state = initial_state
  for i in tf.range(max_seq_len):
    state = rnn_step(input_data[i], state)
    states = states.write(i, state)
  return tf.transpose(states.stack(), [1, 0, 2])

dynamic_rnn(rnn_step,
            tf.random.uniform([batch_size, seq_len, feature_size]),
            tf.zeros([batch_size, feature_size]))

<tf.Tensor: shape=(2, 3, 4), dtype=float32, numpy=
array([[[0.06309307, 0.9938811 , 0.90789986, 0.42136216],
        [0.44997275, 1.9107027 , 1.0716251 , 0.717237  ],
        [0.6026064 , 2.1622117 , 1.4164022 , 1.4153863 ]],

       [[0.04946005, 0.69127274, 0.56848884, 0.22406638],
        [0.8148316 , 1.0278493 , 0.6207781 , 1.1935129 ],
        [0.9178308 , 1.320889  , 0.989761  , 2.0120025 ]]], dtype=float32)>

Ограничения

Function TensorFlow имеет несколько ограничений по дизайну, о которых следует помнить при преобразовании функции Python в Function .

Выполнение побочных эффектов Python

Побочные эффекты, такие как печать, добавление к спискам и изменение глобальных переменных, могут вести себя внутри Function неожиданно, иногда выполняясь дважды или не все. Они происходят только при первом вызове Function с набором входных данных. После этого отслеженный tf.Graph выполняется повторно, без выполнения кода Python.

Общее эмпирическое правило заключается в том, чтобы не полагаться на побочные эффекты Python в своей логике и использовать их только для отладки трассировок. В противном случае API-интерфейсы TensorFlow, такие как tf.data , tf.print , tf.summary , tf.Variable.assign и tf.TensorArray , являются лучшим способом гарантировать, что ваш код будет выполняться средой выполнения TensorFlow при каждом вызове.

@tf.function
def f(x):
  print("Traced with", x)
  tf.print("Executed with", x)

f(1)
f(1)
f(2)

Traced with 1
Executed with 1
Executed with 1
Traced with 2
Executed with 2

Если вы хотите выполнять код Python при каждом вызове Function , tf.py_function — это выходной люк. Недостатком tf.py_function является то, что он не является переносимым или особенно производительным, не может быть сохранен с помощью SavedModel и плохо работает в распределенных (с несколькими GPU, TPU) настройках. Кроме того, поскольку tf.py_function должен быть подключен к графу, он преобразует все входные/выходные данные в тензоры.

Изменение глобальных и свободных переменных Python

Изменение глобальных и свободных переменных Python считается побочным эффектом Python, поэтому это происходит только во время трассировки.

external_list = []

@tf.function
def side_effect(x):
  print('Python side effect')
  external_list.append(x)

side_effect(1)
side_effect(1)
side_effect(1)
# The list append only happened once!
assert len(external_list) == 1

Python side effect

Иногда неожиданное поведение очень трудно заметить. В приведенном ниже примере counter предназначен для защиты приращения переменной. Однако, поскольку это целое число Python, а не объект TensorFlow, его значение фиксируется во время первой трассировки. Когда используется tf.function , assign_add будет безусловно записываться в базовый граф. Поэтому v будет увеличиваться на 1 каждый раз, когда tf.function . Эта проблема распространена среди пользователей, которые пытаются перенести свой код Tensorflow в режиме Grpah на Tensorflow 2 с помощью декораторов tf.function , когда побочные эффекты python ( counter в примере) используются для определения того, какие операции выполнять ( assign_add в примере ). Обычно пользователи понимают это только после того, как увидят подозрительные числовые результаты или значительно более низкую производительность, чем ожидалось (например, если охраняемая операция очень затратна).

class Model(tf.Module):
  def __init__(self):
    self.v = tf.Variable(0)
    self.counter = 0

  @tf.function
  def __call__(self):
    if self.counter == 0:
      # A python side-effect
      self.counter += 1
      self.v.assign_add(1)

    return self.v

m = Model()
for n in range(3):
  print(m().numpy()) # prints 1, 2, 3

1
2
3

Обходной путь для достижения ожидаемого поведения — использование tf.init_scope для подъема операций за пределы графа функции. Это гарантирует, что приращение переменной выполняется только один раз за время трассировки. Следует отметить, что у init_scope есть и другие побочные эффекты, включая очистку потока управления и градиентную ленту. Иногда использование init_scope может стать слишком сложным для реалистичного управления.

class Model(tf.Module):
  def __init__(self):
    self.v = tf.Variable(0)
    self.counter = 0

  @tf.function
  def __call__(self):
    if self.counter == 0:
      # Lifts ops out of function-building graphs
      with tf.init_scope():
        self.counter += 1
        self.v.assign_add(1)

    return self.v

m = Model()
for n in range(3):
  print(m().numpy()) # prints 1, 1, 1

1
1
1

Таким образом, как правило, вам следует избегать изменения объектов Python, таких как целые числа или контейнеры, такие как списки, которые находятся вне Function . Вместо этого используйте аргументы и объекты TF. Например, в разделе «Накопление значений в цикле» есть один пример того, как могут быть реализованы операции, подобные спискам.

В некоторых случаях вы можете захватывать и манипулировать состоянием, если оно является tf.Variable . Вот как обновляются веса моделей Keras при повторных вызовах одной и той же ConcreteFunction .

Использование итераторов и генераторов Python

Многие функции Python, такие как генераторы и итераторы, полагаются на среду выполнения Python для отслеживания состояния. В целом, несмотря на то, что эти конструкции работают должным образом в активном режиме, они являются примерами побочных эффектов Python и поэтому проявляются только во время трассировки.

@tf.function
def buggy_consume_next(iterator):
  tf.print("Value:", next(iterator))

iterator = iter([1, 2, 3])
buggy_consume_next(iterator)
# This reuses the first value from the iterator, rather than consuming the next value.
buggy_consume_next(iterator)
buggy_consume_next(iterator)

Value: 1
Value: 1
Value: 1

Точно так же, как TensorFlow имеет специализированный tf.TensorArray для конструкций списка, у него есть специализированный tf.data.Iterator для итерационных конструкций. Обзор см. в разделе о преобразованиях AutoGraph . Кроме того, API tf.data может помочь реализовать шаблоны генератора:

@tf.function
def good_consume_next(iterator):
  # This is ok, iterator is a tf.data.Iterator
  tf.print("Value:", next(iterator))

ds = tf.data.Dataset.from_tensor_slices([1, 2, 3])
iterator = iter(ds)
good_consume_next(iterator)
good_consume_next(iterator)
good_consume_next(iterator)

Value: 1
Value: 2
Value: 3

Все выходные данные tf.function должны быть возвращаемыми значениями.

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

Например, функция ниже «просачивает» тензор a через глобальную переменную Python x :

x = None

@tf.function
def leaky_function(a):
  global x
  x = a + 1  # Bad - leaks local tensor
  return a + 2

correct_a = leaky_function(tf.constant(1))

print(correct_a.numpy())  # Good - value obtained from function's returns
try:
  x.numpy()  # Bad - tensor leaked from inside the function, cannot be used here
except AttributeError as expected:
  print(expected)

3
'Tensor' object has no attribute 'numpy'

Это верно, даже если утечка значения также возвращается:

@tf.function
def leaky_function(a):
  global x
  x = a + 1  # Bad - leaks local tensor
  return x  # Good - uses local tensor

correct_a = leaky_function(tf.constant(1))

print(correct_a.numpy())  # Good - value obtained from function's returns
try:
  x.numpy()  # Bad - tensor leaked from inside the function, cannot be used here
except AttributeError as expected:
  print(expected)

@tf.function
def captures_leaked_tensor(b):
  b += x  # Bad - `x` is leaked from `leaky_function`
  return b

with assert_raises(TypeError):
  captures_leaked_tensor(tf.constant(2))

2
'Tensor' object has no attribute 'numpy'
Caught expected exception 
  <class 'TypeError'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/566849597.py", line 21, in <module>
    captures_leaked_tensor(tf.constant(2))
TypeError: Originated from a graph execution error.

The graph execution error is detected at a node built at (most recent call last):
>>>  File /usr/lib/python3.7/runpy.py, line 193, in _run_module_as_main
>>>  File /usr/lib/python3.7/runpy.py, line 85, in _run_code
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel_launcher.py, line 16, in <module>
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/traitlets/config/application.py, line 846, in launch_instance
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/kernelapp.py, line 677, in start
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tornado/platform/asyncio.py, line 199, in start
>>>  File /usr/lib/python3.7/asyncio/base_events.py, line 534, in run_forever
>>>  File /usr/lib/python3.7/asyncio/base_events.py, line 1771, in _run_once
>>>  File /usr/lib/python3.7/asyncio/events.py, line 88, in _run
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/kernelbase.py, line 457, in dispatch_queue
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/kernelbase.py, line 446, in process_one
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/kernelbase.py, line 353, in dispatch_shell
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/kernelbase.py, line 648, in execute_request
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/ipkernel.py, line 353, in do_execute
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/ipykernel/zmqshell.py, line 533, in run_cell
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/IPython/core/interactiveshell.py, line 2902, in run_cell
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/IPython/core/interactiveshell.py, line 2947, in _run_cell
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/IPython/core/async_helpers.py, line 68, in _pseudo_sync_runner
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/IPython/core/interactiveshell.py, line 3173, in run_cell_async
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/IPython/core/interactiveshell.py, line 3364, in run_ast_nodes
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/IPython/core/interactiveshell.py, line 3444, in run_code
>>>  File /tmp/ipykernel_26244/566849597.py, line 7, in <module>
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/util/traceback_utils.py, line 150, in error_handler
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py, line 910, in __call__
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py, line 958, in _call
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py, line 781, in _initialize
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py, line 3157, in _get_concrete_function_internal_garbage_collected
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py, line 3557, in _maybe_define_function
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py, line 3402, in _create_graph_function
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py, line 1143, in func_graph_from_py_func
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py, line 672, in wrapped_fn
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py, line 1125, in autograph_handler
>>>  File /tmp/ipykernel_26244/566849597.py, line 4, in leaky_function
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/util/traceback_utils.py, line 150, in error_handler
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/math_ops.py, line 1383, in binary_op_wrapper
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/util/traceback_utils.py, line 150, in error_handler
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py, line 1096, in op_dispatch_handler
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/math_ops.py, line 1737, in _add_dispatch
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/gen_math_ops.py, line 476, in add_v2
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py, line 746, in _apply_op_helper
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py, line 691, in _create_op_internal
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/framework/ops.py, line 3705, in _create_op_internal
>>>  File /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/framework/ops.py, line 2101, in __init__

Error detected in node 'add' defined at: File "/tmp/ipykernel_26244/566849597.py", line 4, in leaky_function

TypeError: tf.Graph captured an external symbolic tensor. The symbolic tensor 'add:0' created by node 'add' is captured by the tf.Graph being executed as an input. But a tf.Graph is not allowed to take symbolic tensors from another graph as its inputs. Make sure all captured inputs of the executing tf.Graph are not symbolic tensors. Use return values, explicit Python locals or TensorFlow collections to access it. Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information.

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

Распространенные способы утечки локальных тензоров также включают изменение внешней коллекции Python или объекта:

class MyClass:

  def __init__(self):
    self.field = None

external_list = []
external_object = MyClass()

def leaky_function():
  a = tf.constant(1)
  external_list.append(a)  # Bad - leaks tensor
  external_object.field = a  # Bad - leaks tensor

Рекурсивные tf.functions не поддерживаются

Рекурсивные Function не поддерживаются и могут вызывать бесконечные циклы. Например,

@tf.function
def recursive_fn(n):
  if n > 0:
    return recursive_fn(n - 1)
  else:
    return 1

with assert_raises(Exception):
  recursive_fn(tf.constant(5))  # Bad - maximum recursion error.

Caught expected exception 
  <class 'Exception'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/2233998312.py", line 9, in <module>
    recursive_fn(tf.constant(5))  # Bad - maximum recursion error.
tensorflow.python.autograph.impl.api.StagingError: in user code:

    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 4, in recursive_fn  *
        return recursive_fn(n - 1)
    File "/tmp/ipykernel_26244/2233998312.py", line 3, in recursive_fn  *
        if n > 0:
    File "/usr/lib/python3.7/abc.py", line 139, in __instancecheck__
        return _abc_instancecheck(cls, instance)

    RecursionError: maximum recursion depth exceeded while calling a Python object

Даже если кажется, что рекурсивная Function работает, функция python будет отслеживаться несколько раз и может повлиять на производительность. Например,

@tf.function
def recursive_fn(n):
  if n > 0:
    print('tracing')
    return recursive_fn(n - 1)
  else:
    return 1

recursive_fn(5)  # Warning - multiple tracings

tracing
tracing
tracing
tracing
tracing
<tf.Tensor: shape=(), dtype=int32, numpy=1>

Известные проблемы

Если ваша Function не оценивает правильно, ошибка может быть объяснена этими известными проблемами, которые планируется исправить в будущем.

В зависимости от глобальных и свободных переменных Python

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

По этой причине вы должны следовать функциональному стилю программирования, который использует аргументы вместо закрытия внешних имен.

@tf.function
def buggy_add():
  return 1 + foo

@tf.function
def recommended_add(foo):
  return 1 + foo

foo = 1
print("Buggy:", buggy_add())
print("Correct:", recommended_add(foo))

Buggy: tf.Tensor(2, shape=(), dtype=int32)
Correct: tf.Tensor(2, shape=(), dtype=int32)

print("Updating the value of `foo` to 100!")
foo = 100
print("Buggy:", buggy_add())  # Did not change!
print("Correct:", recommended_add(foo))

Updating the value of `foo` to 100!
Buggy: tf.Tensor(2, shape=(), dtype=int32)
Correct: tf.Tensor(101, shape=(), dtype=int32)

Другой способ обновить глобальное значение — сделать его tf.Variable и вместо этого использовать метод Variable.assign .

@tf.function
def variable_add():
  return 1 + foo

foo = tf.Variable(1)
print("Variable:", variable_add())

Variable: tf.Tensor(2, shape=(), dtype=int32)

print("Updating the value of `foo` to 100!")
foo.assign(100)
print("Variable:", variable_add())

Updating the value of `foo` to 100!
Variable: tf.Tensor(101, shape=(), dtype=int32)

В зависимости от объектов Python

Рекомендация передавать объекты Python в качестве аргументов в tf.function имеет ряд известных проблем, которые, как ожидается, будут исправлены в будущем. В общем, вы можете полагаться на непротиворечивую трассировку, если вы используете примитив Python или tf.nest совместимую структуру в качестве аргумента или передаете другой экземпляр объекта в Function . Однако Function не будет создавать новую трассировку при передаче того же объекта, а только изменит его атрибуты .

class SimpleModel(tf.Module):
  def __init__(self):
    # These values are *not* tf.Variables.
    self.bias = 0.
    self.weight = 2.

@tf.function
def evaluate(model, x):
  return model.weight * x + model.bias

simple_model = SimpleModel()
x = tf.constant(10.)
print(evaluate(simple_model, x))

tf.Tensor(20.0, shape=(), dtype=float32)

print("Adding bias!")
simple_model.bias += 5.0
print(evaluate(simple_model, x))  # Didn't change :(

Adding bias!
tf.Tensor(20.0, shape=(), dtype=float32)

Использование той же Function для оценки обновленного экземпляра модели приведет к ошибкам, поскольку обновленная модель имеет тот же ключ кэша , что и исходная модель.

По этой причине вам рекомендуется написать свою Function , чтобы избежать зависимости от изменяемых атрибутов объекта или создания новых объектов.

Если это невозможно, одним из обходных путей является создание новых Function каждый раз, когда вы изменяете свой объект, чтобы принудительно отследить:

def evaluate(model, x):
  return model.weight * x + model.bias

new_model = SimpleModel()
evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x)
# Don't pass in `new_model`, `Function` already captured its state during tracing.
print(evaluate_no_bias(x))

tf.Tensor(20.0, shape=(), dtype=float32)

print("Adding bias!")
new_model.bias += 5.0
# Create new Function and ConcreteFunction since you modified new_model.
evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x)
print(evaluate_with_bias(x)) # Don't pass in `new_model`.

Adding bias!
tf.Tensor(25.0, shape=(), dtype=float32)

Поскольку повторная трассировка может быть дорогостоящей , вы можете использовать tf.Variable в качестве атрибутов объекта, которые можно видоизменять (но не изменять, осторожно!) для получения аналогичного эффекта без необходимости повторной трассировки.

class BetterModel:

  def __init__(self):
    self.bias = tf.Variable(0.)
    self.weight = tf.Variable(2.)

@tf.function
def evaluate(model, x):
  return model.weight * x + model.bias

better_model = BetterModel()
print(evaluate(better_model, x))

tf.Tensor(20.0, shape=(), dtype=float32)

print("Adding bias!")
better_model.bias.assign_add(5.0)  # Note: instead of better_model.bias += 5
print(evaluate(better_model, x))  # This works!

Adding bias!
tf.Tensor(25.0, shape=(), dtype=float32)

Создание tf.Variables

Function поддерживает только одноэлементные tf.Variable созданные один раз при первом вызове и повторно используемые при последующих вызовах функций. Приведенный ниже фрагмент кода будет создавать новую tf.Variable при каждом вызове функции, что приводит к исключению ValueError .

Пример:

@tf.function
def f(x):
  v = tf.Variable(1.0)
  return v

with assert_raises(ValueError):
  f(1.0)

Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/3018268426.py", line 7, in <module>
    f(1.0)
ValueError: in user code:

    File "/tmp/ipykernel_26244/3018268426.py", line 3, in f  *
        v = tf.Variable(1.0)

    ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.

Распространенный шаблон, используемый для обхода этого ограничения, заключается в том, чтобы начать со значения Python None, а затем условно создать tf.Variable , если значение равно None:

class Count(tf.Module):
  def __init__(self):
    self.count = None

  @tf.function
  def __call__(self):
    if self.count is None:
      self.count = tf.Variable(0)
    return self.count.assign_add(1)

c = Count()
print(c())
print(c())

tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Использование с несколькими оптимизаторами Keras

Вы можете столкнуться с ValueError: tf.function only supports singleton tf.Variables created on the first call. при использовании более чем одного оптимизатора Keras с tf.function . Эта ошибка возникает из-за того, что оптимизаторы внутренне создают tf.Variables при первом применении градиентов.

opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)
opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)

@tf.function
def train_step(w, x, y, optimizer):
   with tf.GradientTape() as tape:
       L = tf.reduce_sum(tf.square(w*x - y))
   gradients = tape.gradient(L, [w])
   optimizer.apply_gradients(zip(gradients, [w]))

w = tf.Variable(2.)
x = tf.constant([-1.])
y = tf.constant([2.])

train_step(w, x, y, opt1)
print("Calling `train_step` with different optimizer...")
with assert_raises(ValueError):
  train_step(w, x, y, opt2)

Calling `train_step` with different optimizer...
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "/tmp/ipykernel_26244/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_26244/3167358578.py", line 18, in <module>
    train_step(w, x, y, opt2)
ValueError: in user code:

    File "/tmp/ipykernel_26244/3167358578.py", line 9, in train_step  *
        optimizer.apply_gradients(zip(gradients, [w]))
    File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py", line 639, in apply_gradients  **
        self._create_all_weights(var_list)
    File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py", line 828, in _create_all_weights
        _ = self.iterations
    File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py", line 835, in __getattribute__
        return super(OptimizerV2, self).__getattribute__(name)
    File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py", line 995, in iterations
        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
    File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py", line 1202, in add_weight
        aggregation=aggregation)
    File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/engine/base_layer_utils.py", line 129, in make_variable
        shape=variable_shape if variable_shape else None)

    ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.

Если вам нужно изменить оптимизатор во время обучения, обходной путь — создать новую Function для каждого оптимизатора, напрямую вызывая ConcreteFunction .

opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)
opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)

# Not a tf.function.
def train_step(w, x, y, optimizer):
   with tf.GradientTape() as tape:
       L = tf.reduce_sum(tf.square(w*x - y))
   gradients = tape.gradient(L, [w])
   optimizer.apply_gradients(zip(gradients, [w]))

w = tf.Variable(2.)
x = tf.constant([-1.])
y = tf.constant([2.])

# Make a new Function and ConcreteFunction for each optimizer.
train_step_1 = tf.function(train_step).get_concrete_function(w, x, y, opt1)
train_step_2 = tf.function(train_step).get_concrete_function(w, x, y, opt2)
for i in range(10):
  if i % 2 == 0:
    train_step_1(w, x, y) # `opt1` is not used as a parameter. 
  else:
    train_step_2(w, x, y) # `opt2` is not used as a parameter.

Использование с несколькими моделями Keras

Вы также можете столкнуться с ValueError: tf.function only supports singleton tf.Variables created on the first call. при передаче разных экземпляров модели в одну и ту же Function .

Эта ошибка возникает из-за того, что модели Keras (у которых не определена входная форма ) и слои Keras создают tf.Variables при первом их вызове. Возможно, вы пытаетесь инициализировать эти переменные внутри Function , которая уже была вызвана. Чтобы избежать этой ошибки, попробуйте вызвать model.build(input_shape) для инициализации всех весов перед обучением модели.

дальнейшее чтение

Чтобы узнать, как экспортировать и загрузить Function , см. руководство по SavedModel . Чтобы узнать больше об оптимизациях графов, которые выполняются после трассировки, см. руководство по Grappler . Чтобы узнать, как оптимизировать конвейер данных и профилировать модель, см. руководство Profiler .