Esta página foi traduzida pela API Cloud Translation.

Switch to English

Melhor desempenho com tf.function

Ver no TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar caderno

Em TensorFlow 2, execução ansioso é ativado por padrão. A interface do usuário é intuitiva e flexível (executar operações únicas é muito mais fácil e rápido), mas isso pode prejudicar o desempenho e a capacidade de implantação.

Você pode usar tf.function para fazer gráficos de seus programas. É uma ferramenta de transformação que cria gráficos de fluxo de dados independentes de Python a partir de seu código Python. Isso irá ajudá-lo a criar modelos de elevada performance e portáteis, e é necessário para usar SavedModel .

Este guia irá ajudá-lo a conceituar como tf.function funciona sob o capô, para que você possa usá-lo efetivamente.

As principais lições e recomendações são:

Debug no modo ansioso, em seguida, decore com @tf.function .
Não confie nos efeitos colaterais do Python, como mutação de objeto ou acréscimos de lista.
tf.function funciona melhor com ops TensorFlow; As chamadas NumPy e Python são convertidas em constantes.

Configurar

import tensorflow as tf

Defina uma função auxiliar para demonstrar os tipos de erros que você pode encontrar:

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))

Fundamentos

Uso

A Function que você define (por exemplo, através da aplicação do @tf.function decorador) é como uma operação de núcleo TensorFlow: Você pode executá-lo ansiosamente; você pode calcular gradientes; e assim por diante.

@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-19 01:26:31.359247: 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-19 01:26:31.367277: 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-19 01:26:31.368263: 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-19 01:26:31.370408: 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-19 01:26:31.371033: 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-19 01:26:31.372158: 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-19 01:26:31.373202: 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-19 01:26:32.003755: 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-19 01:26:32.004773: 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-19 01:26:32.005698: 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-19 01:26:32.006564: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14648 MB memory:  -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:00:05.0, compute capability: 7.0
2021-08-19 01:26:32.343820: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 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>

Você pode usar Function s dentro de outra Function s.

@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 pode ser mais rápido do que o código ansioso, especialmente para gráficos com muitas pequenas ops. Mas para gráficos com algumas operações caras (como convoluções), você pode não ver muita aceleração.

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-19 01:26:35.135723: I tensorflow/stream_executor/cuda/cuda_dnn.cc:369] Loaded cuDNN version 8100
Eager conv: 0.004266158999996605
Function conv: 0.004805293000003985
Note how there's not much difference in performance for convolutions
2021-08-19 01:26:40.338397: I tensorflow/core/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory

Rastreamento

Esta seção apresenta como Function funciona sob o capô, incluindo detalhes de implementação que podem mudar no futuro. No entanto, uma vez que você entender por que e quando o rastreamento acontece, é muito mais fácil usar tf.function efetivamente!

O que é "rastreamento"?

A Function executa o seu programa em um gráfico TensorFlow . No entanto, um tf.Graph não pode representar todas as coisas que você escreveria em um programa TensorFlow ansioso. Por exemplo, Python suporta polimorfismo, mas tf.Graph requer suas entradas para ter um tipo de dados especificado e dimensão. Ou você pode realizar tarefas secundárias, como ler argumentos de linha de comando, gerar um erro ou trabalhar com um objeto Python mais complexo; nenhuma dessas coisas pode ser executado em um tf.Graph .

Function preenche esta lacuna, separando o seu código em duas etapas:

1) Na primeira fase, chamada de "rastreamento", Function cria uma nova tf.Graph . Código Python funciona normalmente, mas todas as operações TensorFlow (como a adição de dois tensores) são diferidos: eles são capturados pelo tf.Graph e não correr.

2) Na segunda etapa, um tf.Graph que contém tudo o que foi adiada na primeira fase é executado. Este estágio é muito mais rápido do que o estágio de rastreamento.

Dependendo suas entradas, Function nem sempre executar a primeira etapa quando é chamado. Consulte "Regras de rastreamento" abaixo para obter uma melhor noção de como ele faz essa determinação. Pular o primeiro estágio e executar apenas o segundo é o que oferece o alto desempenho do TensorFlow.

Quando Function não decidir traçar, a fase de rastreamento é imediatamente seguido por uma segunda etapa, assim que chamar a Function tanto cria e executa o tf.Graph . Mais tarde você vai ver como você pode executar apenas a fase de rastreamento com get_concrete_function .

Quando passar argumentos de tipos diferentes em uma Function , ambas as fases são executados:

@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)

Note que se você repetidamente chamar uma Function com o mesmo tipo de argumento, TensorFlow vai pular a fase de rastreamento e reutilizar um gráfico previamente traçadas, como o gráfico gerado seria idêntico.

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

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

Você pode usar pretty_printed_concrete_signatures() para ver todos os vestígios disponíveis:

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=()

Até agora, você viu que tf.function cria um cache, camada de expedição dinâmica sobre o gráfico de TensorFlow rastreamento lógica. Para ser mais específico sobre a terminologia:

A tf.Graph é a representação crua, língua-agnóstico, portátil de um cálculo TensorFlow.
A ConcreteFunction envolve um tf.Graph .
A Function gere um cache de ConcreteFunction s e escolhe o caminho certo para suas entradas.
tf.function envolve uma função Python, retornando uma Function objeto.
Traçando cria um tf.Graph e enrola-lo numa ConcreteFunction , também conhecido como um traço.

Regras de rastreamento

A Function determina se reutilizar um traçado ConcreteFunction computando a chave de cache de argumentos de uma entrada e kwargs. A chave de cache é uma chave que identifica um ConcreteFunction com base nos argumentos de entrada e kwargs da Function chamada, de acordo com as seguintes regras (que podem mudar):

A chave gerada por um tf.Tensor é a sua forma e dtipo.
A chave gerada por um tf.Variable é uma id única variável.
A chave gerada por um pitão primitiva (como int , float , str ) é o seu valor.
A chave gerada por nested dict s, list s, tuple s, namedtuple s, e attr s é a tupla achatada de folha-chaves (veja nest.flatten ). (Como resultado desse achatamento, chamar uma função concreta com uma estrutura de aninhamento diferente daquela usada durante o rastreamento resultará em um TypeError).
Para todos os outros tipos de Python, a chave é exclusiva do objeto. Desta forma, uma função ou método é rastreado independentemente para cada instância com a qual é chamado.

Controlando retraçando

Refazendo, que é quando a sua Function cria mais de um traço, ajuda a garantir que TensorFlow gera gráficos corretos para cada conjunto de entradas. No entanto, o rastreamento é uma operação cara! Se a sua Function reconstitui um novo gráfico para cada chamada, você verá que seus executa o código mais lentamente do que se você não usar tf.function .

Para controlar o comportamento de rastreamento, você pode usar as seguintes técnicas:

Especifique input_signature em tf.function para rastreamento limite.

@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_25646/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_25646/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_25646/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_25646/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))

Especificar um [Nenhum] dimensão em tf.TensorSpec para permitir a flexibilidade na reutilização de rastreio.
Desde TensorFlow corresponde tensores com base na sua forma, utilizando um None dimensão como um carácter universal irá permitir Function s aos vestígios de reutilização para entrada de tamanho variável. Entrada de tamanho variável pode ocorrer se você tem as sequências de comprimento diferente, ou imagens de diferentes tamanhos para cada lote (Veja as transformadoras e Sonho profundo tutoriais por exemplo).

@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)

Lance argumentos Python para tensores para reduzir o retrocesso.
Muitas vezes, argumentos Python são utilizados para hiperparâmetros controle e construções gráfico - por exemplo, num_layers=10 ou training=True ou nonlinearity='relu' . Portanto, se o argumento do Python mudar, faz sentido que você tenha que refazer o gráfico.
No entanto, é possível que um argumento Python não esteja sendo usado para controlar a construção do gráfico. Nesses casos, uma mudança no valor do Python pode acionar um retrocesso desnecessário. Tome, por exemplo, este loop de treinamento, que o AutoGraph irá desenrolar dinamicamente. Apesar dos vários traços, o gráfico gerado é realmente idêntico, portanto, o refazer é desnecessário.

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

Se você precisa de força retracing, criar uma nova Function . Separadas Function objetos não são garantidos aos vestígios de acções.

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

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

Tracing!
Executing
Tracing!
Executing

Obtendo funções concretas

Cada vez que uma função é rastreada, uma nova função concreta é criada. Você pode obter directamente uma função concreta, usando 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)

Imprimindo um ConcreteFunction exibe um resumo dos seus argumentos de entrada (com tipos) e seu tipo de saída.

print(double_strings)

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

Você também pode recuperar diretamente a assinatura de uma função concreta.

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

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

Usar um traço concreto com tipos incompatíveis gerará um erro

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_25646/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_25646/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]

Você pode notar que os argumentos do Python recebem um tratamento especial na assinatura de entrada de uma função concreta. Antes do TensorFlow 2.3, os argumentos do Python eram simplesmente removidos da assinatura da função concreta. A partir do TensorFlow 2.3, os argumentos do Python permanecem na assinatura, mas são restritos para assumir o valor definido durante o rastreamento.

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

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

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

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

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

Caught expected exception 
  <class 'TypeError'>:
Traceback (most recent call last):
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1721, in _call_impl
    cancellation_manager)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1766, 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_25646/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_25646/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

Obtenção de gráficos

Cada função de concreto é um invólucro que pode ser chamado em torno de um tf.Graph . Embora recuperar o real tf.Graph objeto não é algo que você normalmente precisa fazer, você pode obtê-lo facilmente a partir de qualquer função concreta.

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

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

Depurando

Em geral, a depuração do código é mais fácil no modo ansioso do que dentro tf.function . Você deve garantir que seu código é executado no modo ansioso livre de erros antes de decorar com tf.function . Para auxiliar no processo de depuração, você pode chamar tf.config.run_functions_eagerly(True) para globalmente desativar e reativar tf.function .

Quando rastrear problemas que só aparecem dentro tf.function , aqui estão algumas dicas:

Velhos Python simples print chamadas apenas executar durante o rastreamento, ajudando a rastrear quando a sua função obtém (re) rastreado.
tf.print chamadas irá executar todas as vezes, e pode ajudá-lo a rastrear valores intermediários durante a execução.
tf.debugging.enable_check_numerics é uma maneira fácil de rastrear onde NaNs e Inf são criados.
pdb (o depurador Python ) pode ajudar você a entender o que está acontecendo durante o rastreamento. (Aviso: pdb vai deixá-lo no código-fonte transformou-autógrafo.)

Transformações de AutoGraph

Autograph é uma biblioteca que é ativado por padrão no tf.function , e transforma um subconjunto de código ansioso Python em ops TensorFlow gráfico compatível. Isso inclui controle de fluxo como if , for , while .

Ops TensorFlow como tf.cond e tf.while_loop continuar a trabalhar, mas o fluxo de controle é muitas vezes mais fácil de escrever e entender quando escrito em 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.680332422 0.0791739225 0.0669257641 0.701964378 0.535618901]
[0.591735482 0.0790088922 0.0668260083 0.605613232 0.489664495]
[0.531142652 0.0788449 0.0667267069 0.541031778 0.453950107]
[0.486254066 0.078681916 0.0666278452 0.493768573 0.425140589]
[0.451238334 0.0785199404 0.0665294155 0.457202375 0.40125224]
[0.422916353 0.0783589631 0.0664314255 0.427801341 0.38101992]
[0.399384439 0.0781989694 0.0663338676 0.403482229 0.363592863]
[0.379422128 0.0780399516 0.0662367418 0.382924527 0.348374784]
[0.362205505 0.0778819 0.0661400482 0.365244567 0.3349334]
[0.347155213 0.0777248144 0.0660437644 0.349825174 0.322946697]
[0.333850235 0.0775686726 0.0659479052 0.336220473 0.312168896]
[0.321976125 0.0774134621 0.0658524558 0.324099 0.302408934]
[0.311292648 0.0772591755 0.0657574236 0.31320855 0.293515563]
[0.301612616 0.0771058127 0.0656628 0.30335319 0.28536731]
[0.292787671 0.0769533589 0.0655685887 0.294378221 0.27786532]
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.28469858, 0.07680181, 0.06547478, 0.28615952, 0.27092823],
      dtype=float32)>

Se você estiver curioso, pode inspecionar o código que o autógrafo gera.

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)

Condicionais

Autograph irá converter alguns if <condition> declarações para o equivalente tf.cond chamadas. Essa substituição é feita se <condition> é um Tensor. Caso contrário, o if instrução é executada como uma condicional Python.

Uma condicional Python é executada durante o rastreamento, então exatamente uma ramificação da condicional será adicionada ao gráfico. Sem o AutoGraph, este gráfico rastreado seria incapaz de tomar o ramo alternativo se houver fluxo de controle dependente de dados.

tf.cond traços e acrescenta ambos os ramos da condicional para o gráfico, selecionar dinamicamente uma filial em tempo de execução. O rastreamento pode ter efeitos colaterais indesejados; confira efeitos rastreamento de autógrafos para mais informações.

@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

Veja a documentação de referência para restrições adicionais sobre Autograph-convertidos se declarações.

rotações

Autograph irá converter alguns for e while declarações para o TensorFlow equivalente looping ops, como tf.while_loop . Se não for convertido, o for ou while laço é executado como um loop de Python.

Esta substituição é feita nas seguintes situações:

for x in y : se y é um Tensor, convertido ao tf.while_loop . No caso especial em que y é um tf.data.Dataset , uma combinação de tf.data.Dataset ops são gerados.
while <condition> : se <condition> é um Tensor, convertido ao tf.while_loop .

A Python executa laço durante o rastreamento, adicionando ops adicionais para o tf.Graph para cada iteração do loop.

Um loop do TensorFlow rastreia o corpo do loop e seleciona dinamicamente quantas iterações serão executadas no tempo de execução. O corpo do laço aparece apenas uma vez no gerado tf.Graph .

Veja a documentação de referência para restrições adicionais sobre Autograph-convertidos for e while declarações.

Loop de dados Python

Um problema comum é de varrer dados Python / Numpy dentro de um tf.function . Este loop será executado durante o processo de rastreamento, adicionando uma cópia do seu modelo para o tf.Graph para cada iteração do loop.

Se você quiser quebrar o ciclo completo de formação em tf.function , a forma mais segura de fazer isso é envolver os seus dados como um tf.data.Dataset para que autógrafo vai dinamicamente desenrolar do ciclo de formação.

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

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

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

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

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

Quando o encapsulamento de informação Python / Numpy num conjunto de dados, estar consciente de tf.data.Dataset.from_generator contra tf.data.Dataset.from_tensors . O primeiro vai manter os dados em Python e buscá-la via tf.py_function que pode ter implicações de desempenho, enquanto o segundo vai incluir uma cópia dos dados como um grande tf.constant() nó no gráfico, o que pode ter implicações de memória.

Leitura de dados de arquivos via TFRecordDataset , CsvDataset , etc. é a maneira mais eficaz para consumir dados, como então TensorFlow si pode gerenciar o carregamento assíncrono e pré-busca de dados, sem ter de envolver Python. Para saber mais, consulte o tf.data : Encanamentos entrada de compilação TensorFlow guiar.

Acumulando valores em um loop

Um padrão comum é acumular valores intermediários de um loop. Normalmente, isso é feito anexando a uma lista Python ou adicionando entradas a um dicionário Python. No entanto, como esses são efeitos colaterais do Python, eles não funcionarão como esperado em um loop desenrolado dinamicamente. Use tf.TensorArray a acumular resultados de um loop dinamicamente desenrolado.

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.9311646 , 0.23201597, 0.27760983, 0.50749147],
        [1.9157217 , 0.73632634, 1.2418742 , 1.3936795 ],
        [2.2432408 , 0.80746305, 1.2652718 , 2.0350993 ]],

       [[0.6350479 , 0.57744336, 0.1134994 , 0.49145043],
        [1.4704436 , 1.2057934 , 0.6300938 , 1.1418717 ],
        [2.2552116 , 1.7597941 , 1.3745583 , 1.9007651 ]]], dtype=float32)>

Limitações

TensorFlow Function tem algumas limitações de design que você deve estar ciente de quando converter uma função Python para uma Function .

Execução de efeitos colaterais do Python

Efeitos colaterais, como a impressão, anexando a listas e mutação globals, podem se comportar de forma inesperada dentro de uma Function , por vezes, a execução de duas ou não todos. Eles só acontecerá a primeira vez que você chamar uma Function com um conjunto de entradas. Posteriormente, o traçado tf.Graph é reexecutado, sem executar o código Python.

A regra geral é evitar depender dos efeitos colaterais do Python em sua lógica e usá-los apenas para depurar seus rastreamentos. Caso contrário, APIs TensorFlow como tf.data , tf.print , tf.summary , tf.Variable.assign e tf.TensorArray são a melhor maneira de garantir o seu código será executado pelo tempo de execução TensorFlow com cada chamada.

@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

Se você gostaria de executar código Python durante cada chamada de uma Function , tf.py_function é uma escotilha de saída. A desvantagem de tf.py_function é que é não portátil ou particularmente alto desempenho, não pode ser salvo com SavedModel, e não funciona bem em configurações distribuídas (multi-GPU, TPU). Também, desde que tf.py_function tem que ser conectado ao gráfico, que lança todas as entradas / saídas para tensores.

Alterar variáveis globais e livres do Python

Alterar Python globais e variáveis livres conta como um efeito colateral Python, por isso só acontece durante o rastreamento.

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

Você deve evitar mutação recipientes como listas, dicionários, outros objetos que vivem fora da Function . Em vez disso, use argumentos e objetos TF. Por exemplo, a seção de "acumular valores em um loop" tem um exemplo de como lista-like operações podem ser implementadas.

Você pode, em alguns casos, a captura e manipular se se trata de um tf.Variable . Isto é como os pesos dos modelos Keras são atualizados com chamadas repetidas para o mesmo ConcreteFunction .

Usando iteradores e geradores Python

Muitos recursos do Python, como geradores e iteradores, dependem do tempo de execução do Python para controlar o estado. Em geral, embora essas construções funcionem conforme o esperado no modo ansioso, elas são exemplos dos efeitos colaterais do Python e, portanto, acontecem apenas durante o rastreamento.

@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

Assim como como TensorFlow tem uma especializada tf.TensorArray para a lista de construções, tem uma especializada tf.data.Iterator para iteração construções. Veja a seção sobre transformações autógrafo para uma visão geral. Além disso, o tf.data API pode ajudar a implementar os padrões de geradores:

@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

Excluindo tf.Variables entre `Function` chamadas

Outro erro que você pode encontrar é uma variável com coleta de lixo. ConcreteFunction s única reter WeakRefs às variáveis que eles perto mais, então você deve manter uma referência para todas as variáveis.

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_25646/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_25646/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_25646/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_25646/3862898592.py:4) ]]
     [[ReadVariableOp/_2]]
0 successful operations.
0 derived errors ignored. [Op:__inference_f_772]

Function call stack:
f -> f

Problemas Conhecidos

Se a sua Function não está avaliando corretamente, o erro pode ser explicada por esses problemas conhecidos que são planejadas para ser corrigido no futuro.

Dependendo das variáveis globais e livres do Python

Function cria uma nova ConcreteFunction quando chamado com um novo valor de um argumento Python. No entanto, ele não faz isso para o encerramento Python, globals, ou nonlocals dessa Function . Se o seu valor muda entre chamadas para a Function , a Function ainda usará os valores que tinham quando foi rastreada. Isso é diferente de como as funções regulares do Python funcionam.

Por esse motivo, recomendamos um estilo de programação funcional que use argumentos em vez de fechar sobre nomes externos.

@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)

Você pode fechar nomes externos, contanto que não atualize seus valores.

Dependendo de objetos Python

A recomendação para passar objetos Python como argumentos em tf.function tem uma série de problemas conhecidos, que são esperados para ser corrigido no futuro. Em geral, você pode contar com o rastreamento consistente se você usar um primitivo ou Python tf.nest estrutura -Compatível como um argumento ou passar em uma instância diferente de um objeto em uma Function . No entanto, Function não irá criar um novo rastreamento quando você passar o mesmo objeto e só mudar seus atributos.

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)

Usando a mesma Function para avaliar a instância atualizada do modelo vai ser buggy desde o modelo atualizado tem a mesma chave de cache como o modelo original.

Por essa razão, recomendamos que você escrever sua Function para evitar dependendo dos atributos do objeto mutáveis ou criar novos objetos.

Se isso não for possível, uma solução alternativa é fazer nova Function de cada vez que você modificar seu objeto a força retracing:

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)

Como retracing pode ser caro , você pode usar tf.Variable s como atributos de objeto, que pode ser mutado (mas não alterados, cuidado!) Para um efeito semelhante sem precisar de um retrace.

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)

Criando tf.Variables

Function suporta apenas singleton tf.Variable s criado uma vez na primeira chamada, e reutilizados em chamadas de função subseqüentes. O trecho de código a seguir criaria um novo tf.Variable em cada chamada de função, o que resulta em um ValueError exceção.

Exemplo:

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

    /tmp/ipykernel_25646/3018268426.py:3 f  *
        v = tf.Variable(1.0)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:268 __call__  **
        return cls._variable_v2_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:262 _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:765 invalid_creator_scope
        "tf.function-decorated function tried to create "

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

Um padrão comum usado para contornar esta limitação é começar com um Nenhum valor Python, em seguida, condicionalmente criar o tf.Variable se o valor for Nenhum:

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)

Usando com vários otimizadores Keras

Você pode encontrar ValueError: tf.function only supports singleton tf.Variables created on the first call. quando se utiliza mais de um otimizador Keras com um tf.function . Este erro ocorre porque otimizadores criar internamente tf.Variables quando se aplicam gradientes pela primeira vez.

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_25646/3551158538.py", line 8, in assert_raises
    yield
  File "/tmp/ipykernel_25646/3167358578.py", line 18, in <module>
    train_step(w, x, y, opt2)
ValueError: in user code:

    /tmp/ipykernel_25646/3167358578.py:9 train_step  *
        optimizer.apply_gradients(zip(gradients, [w]))
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:628 apply_gradients  **
        self._create_all_weights(var_list)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:813 _create_all_weights
        _ = self.iterations
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:820 __getattribute__
        return super(OptimizerV2, self).__getattribute__(name)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:980 iterations
        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:1186 add_weight
        aggregation=aggregation)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:818 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/engine/base_layer_utils.py:129 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:266 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:227 _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:765 invalid_creator_scope
        "tf.function-decorated function tried to create "

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

Se você precisa mudar o otimizador durante o treinamento, a solução é criar uma nova Function para cada otimizador, chamando a ConcreteFunction diretamente.

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.

Usando com vários modelos Keras

Você também pode encontrar ValueError: tf.function only supports singleton tf.Variables created on the first call. quando passar diferentes instâncias do modelo à mesma Function .

Este erro ocorre porque os modelos Keras (que não tem a sua forma de entrada definido ) e camadas Keras criar tf.Variables s quando são chamados em primeiro lugar. Você pode estar tentando inicializar essas variáveis dentro de uma Function , que já foi chamado. Para evitar esse erro, tente chamar model.build(input_shape) para inicializar todos os pesos antes de treinar o modelo.

Leitura adicional

Para saber mais sobre como exportar e carregar uma Function , consulte o guia SavedModel . Para saber mais sobre otimizações de gráficos que são executadas depois de traçar, consulte o guia de Grappler . Para saber como otimizar o seu pipeline de dados e perfil de seu modelo, consulte o guia Profiler .