Лучшая производительность с 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
2021-08-02 22:19:19.897689: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0

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

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.]]
2021-08-02 22:19:21.085343: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcuda.so.1
2021-08-02 22:19:21.729717: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:21.730560: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties: 
pciBusID: 0000:00:05.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2021-08-02 22:19:21.730612: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
2021-08-02 22:19:21.734128: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublas.so.11
2021-08-02 22:19:21.734207: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublasLt.so.11
2021-08-02 22:19:21.735374: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcufft.so.10
2021-08-02 22:19:21.735738: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcurand.so.10
2021-08-02 22:19:21.736864: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcusolver.so.11
2021-08-02 22:19:21.737722: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcusparse.so.11
2021-08-02 22:19:21.737879: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudnn.so.8
2021-08-02 22:19:21.737966: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:21.738897: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:21.739704: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
2021-08-02 22:19:21.740395: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-08-02 22:19:21.740962: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:21.741884: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties: 
pciBusID: 0000:00:05.0 name: Tesla V100-SXM2-16GB computeCapability: 7.0
coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 15.78GiB deviceMemoryBandwidth: 836.37GiB/s
2021-08-02 22:19:21.741958: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:21.742923: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:21.743725: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
2021-08-02 22:19:21.743765: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
2021-08-02 22:19:22.334929: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1258] Device interconnect StreamExecutor with strength 1 edge matrix:
2021-08-02 22:19:22.334965: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1264]      0 
2021-08-02 22:19:22.334973: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1277] 0:   N 
2021-08-02 22:19:22.335203: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:22.336233: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:22.337138: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-02 22:19:22.337941: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1418] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 14646 MB memory) -> physical GPU (device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:00:05.0, compute capability: 7.0)
2021-08-02 22:19:22.644502: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)
2021-08-02 22:19:22.644894: I tensorflow/core/platform/profile_utils/cpu_utils.cc:114] CPU Frequency: 2000194999 Hz
<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]))
2021-08-02 22:19:22.724228: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublas.so.11
2021-08-02 22:19:23.126294: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublasLt.so.11
<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")
2021-08-02 22:19:23.150560: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudnn.so.8
2021-08-02 22:19:23.552403: I tensorflow/stream_executor/cuda/cuda_dnn.cc:359] Loaded cuDNN version 8100
Eager conv: 0.004413468006532639
Function conv: 0.004342082997027319
Note how there's not much difference in performance for convolutions

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

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

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

Function запускает программу в TensorFlow графике . Однако 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 s и выбирает правильную для ваших входов.
  • 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 "/tmp/ipykernel_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/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_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/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 s для повторного использования следов для ввода переменно размера. Вход переменно размера может произойти , если у вас есть последовательности разной длины, или изображений различных размеров для каждой партии (см Transformer и глубокий сон учебники, например).

@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 "/tmp/ipykernel_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/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 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 "/tmp/ipykernel_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/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 , вот несколько советов:

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

Преобразования 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.908887744 0.961447954 0.0590943098 0.0116709471 0.153706789]
[0.720598102 0.744922042 0.059025608 0.0116704153 0.152507633]
[0.617279708 0.632109761 0.0589571483 0.0116698844 0.151336148]
[0.549231172 0.559503257 0.0588889234 0.0116693536 0.150191292]
[0.499943793 0.507608831 0.0588209368 0.0116688227 0.149072066]
[0.462072939 0.46808 0.058753185 0.0116682919 0.147977531]
[0.431772232 0.4366467 0.0586856641 0.011667761 0.146906793]
[0.406801313 0.410861045 0.0586183853 0.0116672302 0.145858988]
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.3857533 , 0.38920352, 0.05855133, 0.0116667 , 0.14483333],
      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)

Условные

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

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

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

@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

Смотрите справочную документацию для дополнительных ограничений на 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 ОПС генерируются.
  • в <condition> tf.while_loop while <condition> : если <condition> является тензором, обращенным в tf.while_loop .

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

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

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

Цикл по данным Python

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

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

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 в Dataset, иметь в виду 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.7595345 , 0.8288754 , 0.07420504, 0.85769975],
        [1.5213574 , 0.9227042 , 0.45271885, 1.3538495 ],
        [1.6639059 , 1.3852738 , 1.2920266 , 1.419443  ]],

       [[0.09766459, 0.5416999 , 0.16834688, 0.9857919 ],
        [0.39503837, 1.428416  , 0.5706631 , 1.376408  ],
        [0.8582125 , 1.5586668 , 1.5131071 , 1.4433464 ]]], dtype=float32)>

Ограничения

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

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

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

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

Вы должны избегать мутирует контейнеров , как списки, dicts, другие объекты , которые живут за пределами 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 для итеративных конструкций. Смотрите раздел преобразований Автограф для обзора. Кроме того , 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 "/tmp/ipykernel_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/3862898592.py", line 16, in <module>
    traced_f(4)
tensorflow.python.framework.errors_impl.FailedPreconditionError: 2 root error(s) found.
  (0) Failed precondition:  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 tmp/ipykernel_17014/3862898592.py:4) ]]
  (1) Failed precondition:  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 tmp/ipykernel_17014/3862898592.py:4) ]]
     [[ReadVariableOp/_2]]
0 successful operations.
0 derived errors ignored. [Op:__inference_f_782]

Function call stack:
f -> f

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

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

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

Function создает новый ConcreteFunction при вызове с новым значением аргумента Python. Тем не менее, это не делает , что для закрытия Python, глобал или nonlocals этой 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 -совместимую структуру в качестве аргумента или передать в другом экземпляре объекта в 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 s как атрибуты объекта, которые могут быть мутировавшие (но не изменились, осторожно!) Для подобного эффекта без необходимости обратного хода.

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 "/tmp/ipykernel_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/3018268426.py", line 7, in <module>
    f(1.0)
ValueError: in user code:

    /tmp/ipykernel_17014/3018268426.py: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. при использовании более чем одного 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_17014/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_17014/3167358578.py", line 18, in <module>
    train_step(w, x, y, opt2)
ValueError: in user code:

    /tmp/ipykernel_17014/3167358578.py: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 (которые не имеют их входной формы определяется ) и Keras слои создают tf.Variables s , когда они первый называется. Возможно , вы пытаетесь инициализировать эти переменные внутри Function , которая уже была вызвана. Чтобы избежать этой ошибки, попробуйте вызвать model.build(input_shape) , чтобы инициализировать все веса перед тренировкой модели.

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

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