Reserve a data! O Google I / O retorna de 18 a 20 de maio Registre-se agora
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

No TensorFlow 2, a execução rápida é ativada por padrão. A interface do usuário é intuitiva e flexível (a execução de operações únicas é muito mais fácil e rápida), 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 o ajudará a criar modelos portáteis e de SavedModel desempenho, e é necessário usar o SavedModel .

Este guia o ajudará a conceituar como tf.function funciona nos bastidores para que você possa usá-lo com eficácia.

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

  • Depure no modo @tf.function e, 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 operações do 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

Uma Function você define (por exemplo, aplicando o decorador @tf.function ) é como uma operação principal do TensorFlow: você pode executá-la avidamente; 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.]]

<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 outras 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ápida do que o código rápido, especialmente para gráficos com muitas operações pequenas. 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")

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

Rastreamento

Esta seção expõe como Function funciona nos bastidores, incluindo detalhes de implementação que podem mudar no futuro . No entanto, depois de entender por que e quando o rastreamento acontece, é muito mais fácil usar tf.function eficaz!

O que é "rastreamento"?

Uma Function executa seu programa em um gráfico do 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 que suas entradas tenham um tipo de dados e dimensão especificados. 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 executada em um tf.Graph .

Function preenche essa lacuna separando seu código em dois estágios:

1) No primeiro estágio, conhecido como " rastreamento ", Function cria um novo tf.Graph . O código Python é executado normalmente, mas todas as operações do TensorFlow (como adicionar dois Tensores) são adiadas : elas são capturadas pelo tf.Graph e não são executadas.

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

Dependendo de suas entradas, Function nem sempre executará o primeiro estágio quando for chamada. Consulte "Regras de rastreamento" abaixo para ter uma ideia melhor de como isso determina. Pular o primeiro estágio e executar apenas o segundo é o que oferece o alto desempenho do TensorFlow.

Quando Function decide rastrear, o estágio de rastreamento é imediatamente seguido pelo segundo estágio, portanto, chamar a Function cria e executa o tf.Graph . Posteriormente, você verá como pode executar apenas o estágio de rastreamento com get_concrete_function .

Quando passamos argumentos de diferentes tipos para uma Function , os dois estágios 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)

Observe que, se você chamar repetidamente uma Function com o mesmo tipo de argumento, o TensorFlow irá pular o estágio de rastreamento e reutilizar um gráfico previamente rastreado, pois 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 rastreamentos disponíveis:

print(double.pretty_printed_concrete_signatures())

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

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

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

Até agora, você viu que tf.function cria uma camada de envio dinâmica em cache sobre a lógica de rastreamento de gráfico do TensorFlow. Para ser mais específico sobre a terminologia:

  • Um tf.Graph é a representação bruta, tf.Graph linguagem e portátil de uma computação do TensorFlow.
  • Um ConcreteFunction envolve um tf.Graph .
  • Uma Function gerencia um cache de ConcreteFunction escolhe o certo para suas entradas.
  • tf.function envolve uma função Python, retornando um objeto Function .
  • O rastreamento cria um tf.Graph e o envolve em um ConcreteFunction , também conhecido como rastreamento.

Regras de rastreamento

Uma Function determina se deve reutilizar um ConcreteFunction rastreado calculando uma chave de cache a partir dos args e kwargs de uma entrada. Uma chave de cache é uma chave que identifica um ConcreteFunction base nos argumentos e kwargs de entrada da chamada de Function , de acordo com as seguintes regras (que podem mudar):

  • A chave gerada para um tf.Tensor é sua forma e tipo.
  • A chave gerada para um tf.Variable é um id de variável exclusivo.
  • A chave gerada para um primitivo Python (como int , float , str ) é seu valor.
  • A chave gerada para dict s aninhados, list s, tuple s, namedtuple s e attr s é a tupla achatada de chaves-folha (consultenest.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

Retracing, que é quando sua Function cria mais de um trace, ajuda a garantir que o TensorFlow gere gráficos corretos para cada conjunto de entradas. No entanto, o rastreamento é uma operação cara! Se sua Function refaz um novo gráfico para cada chamada, você descobrirá que seu código executa mais lentamente do que se você não usasse tf.function .

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

  • Especifique input_signature em tf.function para limitar o rastreio.
@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])))
# We 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]]))

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

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([4 1], shape=(2,), dtype=int32)
Caught expected exception 
  <class 'ValueError'>:
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-20f544b8adbf>", line 9, in <module>
    next_collatz(tf.constant([[1, 2], [3, 4]]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None))
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-20f544b8adbf>", 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))

  • Especifique uma dimensão [Nenhum] em tf.TensorSpec para permitir flexibilidade na reutilização de rastreio.

    Como o TensorFlow combina tensores com base em sua forma, usar uma dimensão None como um caractere curinga permitirá que Function s reutilize traços para entrada de tamanhos variáveis. A entrada de tamanho variável pode ocorrer se você tiver sequências de comprimentos diferentes ou imagens de tamanhos diferentes para cada lote (consulte os tutoriais do Transformer e Deep Dream, 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.

    Freqüentemente, os argumentos Python são usados ​​para controlar hiperparâmetros e construções de gráficos - por exemplo, num_layers=10 ou training=True or 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ê precisar forçar o retrocesso, crie uma nova Function . Objetos de Function separados têm garantia de não compartilhar rastreios.

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

Imprimir um ConcreteFunction exibe um resumo de 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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-e4e2860a4364>", line 2, in <module>
    double_strings(tf.constant(1))
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute __inference_double_162 as input #0(zero-based) was expected to be a string tensor but is a int32 tensor [Op:__inference_double_162]

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.6/site-packages/tensorflow/python/eager/function.py", line 1683, in _call_impl
    cancellation_manager)
  File "/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1728, in _call_with_flat_signature
    self._flat_signature_summary(), ", ".join(sorted(kwargs))))
TypeError: pow(a) got unexpected keyword arguments: b.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-d163f3d206cb>", line 4, in <module>
    square(tf.constant(10.0), b=3)
TypeError: ConcreteFunction pow(a, b) was constructed with int value 2 in b, but was called with int value 3

Obtenção de gráficos

Cada função concreta é um invólucro que pode ser tf.Graph em torno de um tf.Graph . Embora recuperar o objeto tf.Graph real não seja algo que você normalmente precisará 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 de tf.function . Você deve garantir que seu código seja executado sem erros no modo antecipado antes de decorar com tf.function . Para ajudar no processo de depuração, você pode chamar tf.config.run_functions_eagerly(True) para desativar e reativar globalmente tf.function .

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

  • Chamadas de print simples e antigas do Python são executadas apenas durante o rastreamento, ajudando a rastrear quando sua função é (re) rastreada.
  • tf.print chamadas tf.print serão executadas todas as vezes e podem 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 pode ajudá-lo a entender o que está acontecendo durante o rastreamento. (Advertência: o PDB o colocará no código-fonte transformado do AutoGraph.)

Transformações de AutoGraph

AutoGraph é uma biblioteca que está tf.function por padrão em tf.function e transforma um subconjunto de código Python em operações do TensorFlow compatíveis com gráficos. Isso inclui o fluxo de controle como if , for , while .

Operações do TensorFlow, como tf.cond e tf.while_loop continuam funcionando, mas o fluxo de controle costuma ser mais fácil de escrever e entender quando escrito em Python.

# 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.928048491 0.537333608 0.319427252 0.414729953 0.138620138]
[0.729682684 0.490966946 0.308988899 0.392481416 0.137739]
[0.62287122 0.454983532 0.299516946 0.373497456 0.136874482]
[0.553123951 0.425986826 0.290870458 0.357047111 0.13602607]
[0.502857924 0.401961982 0.282935768 0.342610359 0.135193244]
[0.464361787 0.381626487 0.27562 0.329805791 0.134375557]
[0.433632493 0.364119112 0.268846452 0.318346262 0.133572534]
[0.408352554 0.348837078 0.262551099 0.308010817 0.132783771]
[0.387072921 0.335343778 0.256680071 0.298626363 0.132008836]
[0.368834078 0.32331419 0.251187652 0.290055037 0.131247327]
[0.352971435 0.312500536 0.246034727 0.282185435 0.130498841]
[0.339008093 0.302710205 0.241187632 0.274926543 0.129763052]
[0.326591551 0.293790847 0.236617178 0.26820302 0.129039586]
[0.315454811 0.285620153 0.232297987 0.261951953 0.128328085]
[0.305391371 0.278098613 0.228207797 0.256120354 0.127628237]
[0.296238661 0.27114439 0.224326983 0.250663161 0.126939729]
[0.287866682 0.264689356 0.220638305 0.245541915 0.126262262]
[0.280170113 0.25867638 0.217126325 0.240723446 0.12559554]
[0.273062497 0.253057063 0.213777393 0.236178935 0.124939285]
[0.266472191 0.247790173 0.210579231 0.231883332 0.124293216]
[0.260339141 0.242840245 0.207520843 0.227814704 0.12365707]
[0.254612684 0.238176659 0.204592302 0.223953649 0.123030603]
[0.249249727 0.23377277 0.201784685 0.220283121 0.122413576]
[0.244213238 0.229605287 0.199089885 0.216787875 0.12180575]
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.23947136, 0.22565375, 0.19650048, 0.21345437, 0.12120689],
      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 algumas instruções if <condition> nas chamadas tf.cond equivalentes. Esta substituição é feita se a <condition> for um Tensor. Caso contrário, a instrução if é executada como uma condição 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 rastreia e adiciona ambas as ramificações da condicional ao gráfico, selecionando dinamicamente uma ramificação no tempo de execução. O rastreamento pode ter efeitos colaterais indesejados; consulte os efeitos de rastreamento do AutoGraph para obter 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

Consulte a documentação de referência para restrições adicionais sobre instruções if convertidas por AutoGraph.

rotações

O AutoGraph converterá algumas instruções for e while nas operações de loop do TensorFlow equivalentes, 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 for um Tensor, converta para tf.while_loop . No caso especial em que y é umtf.data.Dataset , uma combinação de operaçõestf.data.Dataset é gerada.
  • while <condition> : se <condition> for um Tensor, converta para tf.while_loop .

Um loop Python é executado durante o rastreamento, adicionando ops adicionais ao 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 loop só aparece uma vez no tf.Graph gerado.

Consulte a documentação de referência para obter restrições adicionais sobre as declarações for e while convertidas for AutoGraph.

Loop sobre dados Python

Uma armadilha comum é fazer um loop nos dados Python / Numpy em um tf.function . Este loop será executado durante o processo de rastreamento, adicionando uma cópia do seu modelo ao tf.Graph para cada iteração do loop.

Se você deseja envolver todo o loop de treinamento em tf.function , a maneira mais segura de fazer isso étf.data.Dataset seus dados como umtf.data.Dataset para que o AutoGraph desenrole dinamicamente o loop de treinamento.

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

Ao agrupar dados Python / Numpy em um conjunto de dados, tf.data.Dataset.from_generator atenção a tf.data.Dataset.from_generator versus tf.data.Dataset.from_tensors . O primeiro manterá os dados em Python e os buscará por meio de tf.py_function que pode ter implicações de desempenho, enquanto o último tf.constant() uma cópia dos dados como um grande nó tf.constant() no gráfico, que pode ter implicações de memória.

Lendo dados de arquivos via TFRecordDataset / CsvDataset / etc. é a maneira mais eficaz de consumir dados, pois o próprio TensorFlow pode gerenciar o carregamento assíncrono e a pré-busca de dados, sem precisar envolver o Python. Para saber mais, consulte o guia tf.data .

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 para acumular resultados de um loop desenrolado dinamicamente.

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.8216245 , 0.29562855, 0.379112  , 0.49940717],
        [1.6473945 , 1.039927  , 1.3268942 , 0.5298227 ],
        [2.4393063 , 1.1283967 , 2.087479  , 1.2748951 ]],

       [[0.08016336, 0.73864746, 0.33738315, 0.4542967 ],
        [0.7459605 , 1.307698  , 1.1588445 , 0.9293362 ],
        [1.3752056 , 1.6133544 , 1.8199729 , 1.7356051 ]]], dtype=float32)>

Limitações

A Function TensorFlow tem algumas limitações de design que você deve conhecer ao converter uma função Python em uma Function .

Execução de efeitos colaterais do Python

Os efeitos colaterais, como impressão, acréscimo a listas e alterações globais, podem se comportar inesperadamente dentro de uma Function , às vezes executando duas vezes ou não todas. Eles só acontecem na primeira vez que você chama uma Function com um conjunto de entradas. Posteriormente, o tf.Graph rastreado é 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, as APIs do TensorFlow como tf.data , tf.print , tf.summary , tf.Variable.assign e tf.TensorArray são a melhor maneira de garantir que seu código seja executado pelo tempo de execução do TensorFlow a 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 o código Python durante cada invocação de uma Function , tf.py_function é uma saída de emergência. A desvantagem de tf.py_function é que ele não é portátil ou tem desempenho especial, não pode ser salvo com SavedModel e não funciona bem em configurações distribuídas (multi-GPU, TPU). Além disso, como tf.py_function deve ser conectada ao gráfico, ela converte todas as entradas / saídas em tensores.

Alterar variáveis ​​globais e livres do Python

Alterar as variáveis globais e livres do Python conta como um efeito colateral do Python, portanto, 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 a mutação de containers como listas, dicts e outros objetos que vivem fora da Function . Em vez disso, use argumentos e objetos TF. Por exemplo, a seção "Acumulando valores em um loop" tem um exemplo de como as operações do tipo lista podem ser implementadas.

Você pode, em alguns casos, capturar e manipular o estado se for uma tf.Variable . É assim que 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 o TensorFlow tem um tf.TensorArray especializado para construções de lista, ele tem um tf.data.Iterator especializado para construções de iteração. Consulte a seção sobre Transformações de AutoGraph para uma visão geral. Além disso, a API tf.data pode ajudar a implementar padrões de gerador:

@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 chamadas de Function

Outro erro que você pode encontrar é uma variável com coleta de lixo. ConcreteFunction s apenas retém WeakRefs para as variáveis ​​sobre as quais eles fecham, então você deve reter uma referência para quaisquer 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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-9a93d2e07632>", line 16, in <module>
    traced_f(4)
tensorflow.python.framework.errors_impl.FailedPreconditionError: 2 root error(s) found.
  (0) Failed precondition:  Error while reading resource variable _AnonymousVar3 from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
     [[node ReadVariableOp (defined at <ipython-input-1-9a93d2e07632>:4) ]]
  (1) Failed precondition:  Error while reading resource variable _AnonymousVar3 from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
     [[node ReadVariableOp (defined at <ipython-input-1-9a93d2e07632>:4) ]]
     [[ReadVariableOp/_2]]
0 successful operations.
0 derived errors ignored. [Op:__inference_f_782]

Function call stack:
f -> f

Problemas Conhecidos

Se sua Function não estiver avaliando corretamente, o erro pode ser explicado por esses problemas conhecidos que estão planejados para serem corrigidos no futuro.

Dependendo das variáveis ​​globais e livres do Python

Function cria uma nova ConcreteFunction quando chamada com um novo valor de um argumento Python. No entanto, não faz isso para o fechamento Python, globais ou não locais dessa Function . Se o seu valor mudar entre as chamadas para a Function , a Function ainda usará os valores que tinha 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 de passar objetos Python como argumentos para tf.function tem vários problemas conhecidos, que devem ser corrigidos no futuro. Em geral, você pode contar com um rastreamento consistente se usar uma estrutura primitiva Python ou tf.nest estrutura compatível com tf.nest como um argumento ou passar uma instância diferente de um objeto em uma Function . Porém, Function não irá criar um novo traço quando você passar o mesmo objeto e apenas alterar 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)

Usar a mesma Function para avaliar a instância atualizada do modelo terá muitos bugs, pois o modelo atualizado tem a mesma chave de cache do modelo original.

Por esse motivo, recomendamos que você escreva sua Function para evitar depender de atributos de objetos mutáveis ​​ou crie novos objetos.

Se isso não for possível, uma solução alternativa é fazer novas Function s cada vez que você modificar seu objeto para forçar o retrocesso:

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 retraçar pode ser caro , você pode usar tf.Variable s como atributos de objeto, que podem ser mutados (mas não alterados, cuidado!) Para um efeito semelhante sem a necessidade de retraçar.

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 só suporta a criação de variáveis ​​uma vez, quando chamada pela primeira vez, e depois reutilizá-las. Você não pode criar tf.Variables em novos rastreios. A criação de novas variáveis ​​em chamadas subsequentes não é permitida no momento, mas será no futuro.

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

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

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

Você pode criar variáveis ​​dentro de uma Function , desde que essas variáveis ​​sejam criadas apenas na primeira vez que a função é executada.

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-decorated function tried to create variables on non-first call. ao usar mais de um otimizador Keras com um tf.function . Esse erro ocorre porque os otimizadores criam internamente tf.Variables quando 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 "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-d3d3937dbf1a>", line 18, in <module>
    train_step(w, x, y, opt2)
ValueError: in user code:

    <ipython-input-1-d3d3937dbf1a>:9 train_step  *
        optimizer.apply_gradients(zip(gradients, [w]))
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:604 apply_gradients  **
        self._create_all_weights(var_list)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:781 _create_all_weights
        _ = self.iterations
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:788 __getattribute__
        return super(OptimizerV2, self).__getattribute__(name)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:926 iterations
        aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:1132 add_weight
        aggregation=aggregation)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py:810 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:142 make_variable
        shape=variable_shape if variable_shape else None)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:260 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py:731 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ê precisar alterar o otimizador durante o treinamento, uma solução alternativa é criar uma nova Function para cada otimizador, chamando o 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-decorated function tried to create variables on non-first call. ao passar diferentes instâncias de modelo para a mesma Function .

Esse erro ocorre porque os modelos Keras (que não têm sua forma de entrada definida ) e as camadas Keras criam tf.Variables s quando são chamadas pela primeira vez. Você pode estar tentando inicializar essas variáveis ​​dentro de uma Function , que já foi chamada. Para evitar esse erro, tente chamar model.build(input_shape) para inicializar todos os pesos antes de treinar o modelo.

Leitura adicional

Para saber como exportar e carregar uma Function , consulte o guia SavedModel . Para saber mais sobre as otimizações de gráfico que são realizadas após o traçado, consulte o guia do Grappler . Para saber como otimizar seu pipeline de dados e criar o perfil de seu modelo, consulte o guia do Profiler .