Есть вопрос? Присоединяйтесь к сообществу на форуме TensorFlow. Посетите форум.

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

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

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

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

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

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

  • Отлаживайте в нетерпеливом режиме, а затем украсьте @tf.function .
  • Не полагайтесь на побочные эффекты Python, такие как мутация объекта или добавление списка.
  • tf.function лучше всего работает с TensorFlow ops; Вызовы 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 выполнить ее; вы можете вычислять градиенты; и так далее.

@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 s может быть быстрее, чем нетерпеливый код, особенно для графиков с большим количеством мелких операций. Но для графиков с несколькими дорогостоящими операциями (например, свертки) вы можете не увидеть большого ускорения.

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.23302616399996623
Function conv: 0.21780993200013654
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 (например, добавление двух Tensor) откладываются : они захватываются 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 - это необработанное, не tf.Graph от языка, переносимое представление вычислений TensorFlow.
  • ConcreteFunction оборачивает tf.Graph .
  • Function управляет кешем ConcreteFunction и выбирает правильный вариант для ваших входных данных.
  • tf.function оборачивает функцию Python, возвращая объект Function .
  • Трассировка создает tf.Graph и оборачивает его в ConcreteFunction , также известную как трассировка.

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

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

  • Ключ, сгенерированный для tf.Tensor - это его форма и dtype.
  • Ключ, сгенерированный для tf.Variable является уникальным идентификатором переменной.
  • Ключ, сгенерированный для примитива Python (например, int , float , str ), - это его значение.
  • Ключ, сгенерированный для вложенных dict s, list s, tuple s, namedtuple s и attr s, является уплощенным кортежем листовых ключей (см.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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-14ebce7b7ee8>", 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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-14ebce7b7ee8>", 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")))
Tracing with Tensor("a:0", shape=(), dtype=string)
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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-e4e2860a4364>", 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 1725, in _call_impl
    cancellation_manager)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1770, in _call_with_flat_signature
    self._flat_signature_summary(), ", ".join(sorted(kwargs))))
TypeError: pow(a) got unexpected keyword arguments: b.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-d163f3d206cb>", 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

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

tf.cond tf.while_loop такие как 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.710546374 0.327660799 0.393230557 0.545059443 0.666661739]
[0.611019373 0.316417336 0.374141902 0.496808201 0.582779706]
[0.54484427 0.306263864 0.357609242 0.45960331 0.52468282]
[0.496646136 0.297034383 0.343106419 0.429760844 0.481306106]
[0.459475428 0.288596332 0.330247819 0.405121386 0.44728902]
[0.429656595 0.280842364 0.318743408 0.384322464 0.419668049]
[0.405034214 0.273684502 0.308370233 0.366455346 0.396650732]
[0.384248167 0.267049909 0.298953682 0.350887358 0.377079546]
[0.366391063 0.260877609 0.290354759 0.337162226 0.360168517]
[0.350830942 0.255116194 0.282461286 0.324941576 0.345362455]
[0.337112248 0.2497219 0.275181532 0.313968241 0.332256317]
[0.324896872 0.244657204 0.268439621 0.304042816 0.320546716]
[0.313927948 0.239889801 0.262172282 0.295007944 0.310001194]
[0.304006279 0.235391632 0.256326199 0.286737591 0.300438195]
[0.294974595 0.231138244 0.250856102 0.279129326 0.291713566]
[0.286706954 0.227108166 0.245723218 0.272099048 0.283711195]
[0.279101074 0.223282441 0.240894228 0.265576899 0.276336372]
[0.272072881 0.219644368 0.23634018 0.259504348 0.269510925]
[0.26555258 0.216179058 0.232035935 0.253831863 0.263169676]
[0.259481668 0.212873235 0.227959365 0.24851726 0.257257849]
[0.253810644 0.209715009 0.224091038 0.243524343 0.251728892]
[0.248497337 0.206693679 0.220413819 0.238821834 0.246543139]
[0.243505597 0.20379965 0.216912434 0.2343826 0.241666421]
[0.238804176 0.20102416 0.213573262 0.230182901 0.237069115]
[0.234365895 0.198359385 0.210384145 0.226201892 0.232725471]
[0.230167076 0.195798069 0.207334161 0.222421184 0.228612974]
[0.226186857 0.19333373 0.204413429 0.218824506 0.224711776]
[0.222406894 0.190960392 0.201613098 0.215397388 0.221004337]
[0.218810901 0.188672557 0.198925063 0.212126866 0.217475086]
[0.215384394 0.186465234 0.196342021 0.209001362 0.214110211]
[0.212114424 0.184333771 0.193857312 0.206010461 0.210897282]
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.20898944, 0.18227392, 0.19146483, 0.2031447 , 0.20782518],
      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 оператор if выполняется как условное выражение Python.

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

tf.cond отслеживает и добавляет обе ветви условного 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 в эквивалентные операции цикла tf.while_loop , например tf.while_loop . tf.while_loop . Если не конвертируются, то for или в while цикл выполняется в виде петли Python.

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

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

Цикл Python выполняется во время трассировки, добавляя дополнительные tf.Graph в 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 10 nodes in its graph
train(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 10 nodes in its graph

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

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

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

Распространенным шаблоном является накопление промежуточных значений из цикла. Обычно это достигается путем добавления в список 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.60458577, 0.3308612 , 0.7878152 , 0.3223114 ],
        [0.9110272 , 1.0819752 , 1.7657743 , 1.2409766 ],
        [1.7235098 , 1.5416101 , 2.2929285 , 1.9181627 ]],

       [[0.89487076, 0.22811687, 0.342862  , 0.5752872 ],
        [1.0133923 , 0.28650808, 0.9558767 , 1.0829899 ],
        [1.9280962 , 1.1437279 , 0.9857702 , 1.4834155 ]]], dtype=float32)>

Ограничения

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

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

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

Общее практическое правило - не полагаться на побочные эффекты Python в своей логике и использовать их только для отладки трассировок. В противном случае API-интерфейсы tf.data , tf.print как 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 и плохо работает в распределенных (multi-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

Вам следует избегать изменения контейнеров, таких как списки, словари и другие объекты, которые находятся вне 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 для обзора. Кроме того, tf.data API может помочь реализовать шаблоны генераторов:

@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.Variables между вызовами Function

Другая ошибка, с которой вы можете столкнуться, - это переменная, собранная сборщиком мусора. ConcreteFunction сохраняет WeakRefs только для переменных, которые они закрывают, поэтому вы должны сохранить ссылку на любые переменные.

external_var = tf.Variable(3)
@tf.function
def f(x):
  return x * external_var

traced_f = f.get_concrete_function(4)
print("Calling concrete function...")
print(traced_f(4))

# The original variable object gets garbage collected, since there are no more
# references to it.
external_var = tf.Variable(4)
print()
print("Calling concrete function after garbage collecting its closed Variable...")
with assert_raises(tf.errors.FailedPreconditionError):
  traced_f(4)
Calling concrete function...
tf.Tensor(12, shape=(), dtype=int32)

Calling concrete function after garbage collecting its closed Variable...
Caught expected exception 
  <class 'tensorflow.python.framework.errors_impl.FailedPreconditionError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-9a93d2e07632>", line 16, in <module>
    traced_f(4)
tensorflow.python.framework.errors_impl.FailedPreconditionError:  Could not find variable _AnonymousVar3. This could mean that the variable has been deleted. In TF1, it can also mean the variable is uninitialized. Debug info: container=localhost, status=Not found: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
     [[node ReadVariableOp (defined at <ipython-input-1-9a93d2e07632>:4) ]] [Op:__inference_f_782]

Function call stack:
f

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

Если ваша 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)

Вы можете закрыть внешние имена, если не обновляете их значения.

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

Рекомендация передавать объекты Python в качестве аргументов в tf.function имеет ряд известных проблем, которые, как ожидается, будут исправлены в будущем. В общем, вы можете положиться на последовательную трассировку, если используете в качестве аргумента примитив Python или tf.nest совместимую с 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 качестве атрибутов объекта, которые можно изменять (но не изменять, осторожно!) Для получения аналогичного эффекта без необходимости 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.Variables в новых трассировках. Создание новых переменных в последующих вызовах в настоящее время запрещено, но будет в будущем.

Пример:

@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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-8a0913e250e0>", line 7, in <module>
    f(1.0)
ValueError: in user code:

    <ipython-input-1-8a0913e250e0>:3 f  *
        v = tf.Variable(1.0)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:262 __call__  **
        return cls._variable_v2_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:256 _variable_v2_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:769 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

Вы можете создавать переменные внутри Function если эти переменные создаются только при первом выполнении функции.

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-decorated function tried to create variables on non-first call. при использовании более одного оптимизатора tf.function с tf.function . Эта ошибка возникает из-за того, что оптимизаторы внутренне создают tf.Variables при 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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-d3d3937dbf1a>", line 18, in <module>
    train_step(w, x, y, opt2)
ValueError: in user code:

    <ipython-input-1-d3d3937dbf1a>:9 train_step  *
        optimizer.apply_gradients(zip(gradients, [w]))
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:636 apply_gradients  **
        self._create_all_weights(var_list)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:821 _create_all_weights
        _ = self.iterations
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:828 __getattribute__
        return super(OptimizerV2, self).__getattribute__(name)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:988 iterations
        aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:1194 add_weight
        aggregation=aggregation)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:815 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:139 make_variable
        shape=variable_shape if variable_shape else None)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:260 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:769 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

Если вам нужно изменить оптимизатор во время обучения, обходной путь - создать новую 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-decorated function tried to create variables on non-first call. при передаче разных экземпляров модели в одну и ту же Function .

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

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

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