O Dia da Comunidade de ML é dia 9 de novembro! Junte-nos para atualização de TensorFlow, JAX, e mais Saiba mais

Introdução aos gráficos e tf.function

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Visão geral

Este guia vai além do TensorFlow e Keras para demonstrar como o TensorFlow funciona. Se você quer em vez de se imediatamente começou com Keras, visite a coleção de guias Keras .

Neste guia, você aprenderá como o TensorFlow permite fazer alterações simples em seu código para obter gráficos, como os gráficos são armazenados e representados e como você pode usá-los para acelerar seus modelos.

Esta é uma visão grande-retrato que trata de como tf.function permite-lhe mudar de execução ansioso para execução gráfico. Para uma especificação mais completa de tf.function , vá para o tf.function guia .

O que são gráficos?

Nos três guias anteriores, você correu TensorFlow ansiosamente. Isso significa que as operações do TensorFlow são executadas pelo Python, operação por operação, e retornam os resultados de volta ao Python.

Embora a execução rápida tenha várias vantagens exclusivas, a execução de gráfico permite portabilidade fora do Python e tende a oferecer melhor desempenho. Meios de execução gráfico que computações tensores são executados como um gráfico TensorFlow, por vezes referido como um tf.Graph ou simplesmente um "gráfico".

Os gráficos são estruturas de dados que contenham um conjunto de tf.Operation objectos, que representam as unidades de computação; e tf.Tensor objetos, que representam as unidades de dados que fluem entre operações. Eles são definidos em um tf.Graph contexto. Como esses gráficos são estruturas de dados, eles podem ser salvos, executados e restaurados sem o código Python original.

É assim que um gráfico do TensorFlow que representa uma rede neural de duas camadas se parece quando visualizado no TensorBoard.

Um gráfico simples do TensorFlow

Os benefícios dos gráficos

Com um gráfico, você tem bastante flexibilidade. Você pode usar seu gráfico TensorFlow em ambientes que não têm um interpretador Python, como aplicativos móveis, dispositivos incorporados e servidores de back-end. TensorFlow usa gráficos como o formato para modelos salvos quando exporta-los de Python.

Os gráficos também são facilmente otimizados, permitindo que o compilador faça transformações como:

  • Estaticamente inferir o valor de tensores dobrando nós constantes em sua computação ( "folding constante").
  • Separe as subpartes de um cálculo que são independentes e divida-as entre threads ou dispositivos.
  • Simplifique as operações aritméticas eliminando subexpressões comuns.

Há um sistema de otimização inteira, Grappler , para realizar este e outros aumentos de velocidade.

Em suma, os gráficos são extremamente úteis e deixe o seu TensorFlow correr rápido, em paralelo, e executar de forma eficiente em vários dispositivos.

No entanto, você ainda deseja definir seus modelos de aprendizado de máquina (ou outros cálculos) em Python por conveniência e, em seguida, construir gráficos automaticamente quando precisar deles.

Configurar

import tensorflow as tf
import timeit
from datetime import datetime

Aproveitando os gráficos

Você criar e executar um gráfico em TensorFlow usando tf.function , quer como uma chamada direta ou como um decorador. tf.function leva uma função regular como entrada e retorna a Function . Uma Function é exigível um pitão que constrói gráficos TensorFlow da função Python. Você usa uma Function da mesma forma como o seu equivalente Python.

# Define a Python function.
def a_regular_function(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# `a_function_that_uses_a_graph` is a TensorFlow `Function`.
a_function_that_uses_a_graph = tf.function(a_regular_function)

# Make some tensors.
x1 = tf.constant([[1.0, 2.0]])
y1 = tf.constant([[2.0], [3.0]])
b1 = tf.constant(4.0)

orig_value = a_regular_function(x1, y1, b1).numpy()
# Call a `Function` like a Python function.
tf_function_value = a_function_that_uses_a_graph(x1, y1, b1).numpy()
assert(orig_value == tf_function_value)

No exterior, uma Function se parece com uma função regular que você escreve usando operações TensorFlow. Debaixo , no entanto, é muito diferente. A Function encapsula vários tf.Graph s atrás de uma API . Isso é como Function é capaz de dar-lhe os benefícios da execução gráfico , como a velocidade ea capacidade de colocação.

tf.function se aplica a uma função e todas as outras funções que chama:

def inner_function(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# Use the decorator to make `outer_function` a `Function`.
@tf.function
def outer_function(x):
  y = tf.constant([[2.0], [3.0]])
  b = tf.constant(4.0)

  return inner_function(x, y, b)

# Note that the callable will create a graph that
# includes `inner_function` as well as `outer_function`.
outer_function(tf.constant([[1.0, 2.0]])).numpy()
array([[12.]], dtype=float32)

Se você já usou TensorFlow 1.x, você vai notar que em nenhum momento você precisa definir um Placeholder ou tf.Session .

Converter funções Python em gráficos

Qualquer função que você escreve com TensorFlow irá conter uma mistura de built-in operações TF e lógica Python, tais como if-then cláusulas, loops, break , return , continue , e muito mais. Enquanto as operações TensorFlow são facilmente capturado por um tf.Graph , necessidades lógicas Python específicos para se submeter a um passo extra, a fim de tornar-se parte do gráfico. tf.function utiliza uma biblioteca chamada Autograph ( tf.autograph ) para converter o código Python em código de geração de gráfico.

def simple_relu(x):
  if tf.greater(x, 0):
    return x
  else:
    return 0

# `tf_simple_relu` is a TensorFlow `Function` that wraps `simple_relu`.
tf_simple_relu = tf.function(simple_relu)

print("First branch, with graph:", tf_simple_relu(tf.constant(1)).numpy())
print("Second branch, with graph:", tf_simple_relu(tf.constant(-1)).numpy())
First branch, with graph: 1
Second branch, with graph: 0

Embora seja improvável que você precise visualizar os gráficos diretamente, você pode inspecionar as saídas para verificar os resultados exatos. Não são fáceis de ler, por isso não é preciso olhar muito atentamente!

# This is the graph-generating output of AutoGraph.
print(tf.autograph.to_code(simple_relu))
def tf__simple_relu(x):
    with ag__.FunctionScope('simple_relu', '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 (do_return, retval_)

        def set_state(vars_):
            nonlocal retval_, do_return
            (do_return, retval_) = vars_

        def if_body():
            nonlocal retval_, do_return
            try:
                do_return = True
                retval_ = ag__.ld(x)
            except:
                do_return = False
                raise

        def else_body():
            nonlocal retval_, do_return
            try:
                do_return = True
                retval_ = 0
            except:
                do_return = False
                raise
        ag__.if_stmt(ag__.converted_call(ag__.ld(tf).greater, (ag__.ld(x), 0), None, fscope), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
        return fscope.ret(retval_, do_return)
# This is the graph itself.
print(tf_simple_relu.get_concrete_function(tf.constant(1)).graph.as_graph_def())
node {
  name: "x"
  op: "Placeholder"
  attr {
    key: "_user_specified_name"
    value {
      s: "x"
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
}
node {
  name: "Greater/y"
  op: "Const"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_INT32
        tensor_shape {
        }
        int_val: 0
      }
    }
  }
}
node {
  name: "Greater"
  op: "Greater"
  input: "x"
  input: "Greater/y"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "cond"
  op: "StatelessIf"
  input: "Greater"
  input: "x"
  attr {
    key: "Tcond"
    value {
      type: DT_BOOL
    }
  }
  attr {
    key: "Tin"
    value {
      list {
        type: DT_INT32
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_BOOL
        type: DT_INT32
      }
    }
  }
  attr {
    key: "_lower_using_switch_merge"
    value {
      b: true
    }
  }
  attr {
    key: "_read_only_resource_inputs"
    value {
      list {
      }
    }
  }
  attr {
    key: "else_branch"
    value {
      func {
        name: "cond_false_34"
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
        }
        shape {
        }
      }
    }
  }
  attr {
    key: "then_branch"
    value {
      func {
        name: "cond_true_33"
      }
    }
  }
}
node {
  name: "cond/Identity"
  op: "Identity"
  input: "cond"
  attr {
    key: "T"
    value {
      type: DT_BOOL
    }
  }
}
node {
  name: "cond/Identity_1"
  op: "Identity"
  input: "cond:1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "Identity"
  op: "Identity"
  input: "cond/Identity_1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
library {
  function {
    signature {
      name: "cond_false_34"
      input_arg {
        name: "cond_placeholder"
        type: DT_INT32
      }
      output_arg {
        name: "cond_identity"
        type: DT_BOOL
      }
      output_arg {
        name: "cond_identity_1"
        type: DT_INT32
      }
    }
    node_def {
      name: "cond/Const"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Const"
      }
    }
    node_def {
      name: "cond/Const_1"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Const_1"
      }
    }
    node_def {
      name: "cond/Const_2"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_INT32
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_INT32
            tensor_shape {
            }
            int_val: 0
          }
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Const_2"
      }
    }
    node_def {
      name: "cond/Const_3"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Const_3"
      }
    }
    node_def {
      name: "cond/Identity"
      op: "Identity"
      input: "cond/Const_3:output:0"
      attr {
        key: "T"
        value {
          type: DT_BOOL
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Identity"
      }
    }
    node_def {
      name: "cond/Const_4"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_INT32
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_INT32
            tensor_shape {
            }
            int_val: 0
          }
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Const_4"
      }
    }
    node_def {
      name: "cond/Identity_1"
      op: "Identity"
      input: "cond/Const_4:output:0"
      attr {
        key: "T"
        value {
          type: DT_INT32
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Identity_1"
      }
    }
    ret {
      key: "cond_identity"
      value: "cond/Identity:output:0"
    }
    ret {
      key: "cond_identity_1"
      value: "cond/Identity_1:output:0"
    }
    attr {
      key: "_construction_context"
      value {
        s: "kEagerRuntime"
      }
    }
    arg_attr {
      key: 0
      value {
        attr {
          key: "_output_shapes"
          value {
            list {
              shape {
              }
            }
          }
        }
      }
    }
  }
  function {
    signature {
      name: "cond_true_33"
      input_arg {
        name: "cond_identity_1_x"
        type: DT_INT32
      }
      output_arg {
        name: "cond_identity"
        type: DT_BOOL
      }
      output_arg {
        name: "cond_identity_1"
        type: DT_INT32
      }
    }
    node_def {
      name: "cond/Const"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Const"
      }
    }
    node_def {
      name: "cond/Identity"
      op: "Identity"
      input: "cond/Const:output:0"
      attr {
        key: "T"
        value {
          type: DT_BOOL
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Identity"
      }
    }
    node_def {
      name: "cond/Identity_1"
      op: "Identity"
      input: "cond_identity_1_x"
      attr {
        key: "T"
        value {
          type: DT_INT32
        }
      }
      experimental_debug_info {
        original_node_names: "cond/Identity_1"
      }
    }
    ret {
      key: "cond_identity"
      value: "cond/Identity:output:0"
    }
    ret {
      key: "cond_identity_1"
      value: "cond/Identity_1:output:0"
    }
    attr {
      key: "_construction_context"
      value {
        s: "kEagerRuntime"
      }
    }
    arg_attr {
      key: 0
      value {
        attr {
          key: "_output_shapes"
          value {
            list {
              shape {
              }
            }
          }
        }
      }
    }
  }
}
versions {
  producer: 808
  min_consumer: 12
}

Na maioria das vezes, tf.function funcionará sem considerações especiais. No entanto, existem algumas ressalvas, eo guia tf.function pode ajudar aqui, bem como a referência de autógrafos completa

Polimorfismo: uma Function , muitos gráficos

Um tf.Graph é especializada para um tipo específico de entradas (por exemplo, com um tensores específico dtype ou objectos com o mesmo id() ).

Cada vez que você chamar uma Function com novos dtypes e formas em seus argumentos, Function cria uma nova tf.Graph para os novos argumentos. Os dtypes e formas de uma tf.Graph inputs 's são conhecidos como uma assinatura de entrada ou apenas uma assinatura.

Os Function armazena o tf.Graph correspondente a essa assinatura em um ConcreteFunction . A ConcreteFunction é um wrapper em torno de um tf.Graph .

@tf.function
def my_relu(x):
  return tf.maximum(0., x)

# `my_relu` creates new graphs as it observes more signatures.
print(my_relu(tf.constant(5.5)))
print(my_relu([1, -1]))
print(my_relu(tf.constant([3., -3.])))
tf.Tensor(5.5, shape=(), dtype=float32)
tf.Tensor([1. 0.], shape=(2,), dtype=float32)
tf.Tensor([3. 0.], shape=(2,), dtype=float32)

Se a Function já foi chamado com essa assinatura, Function não cria uma nova tf.Graph .

# These two calls do *not* create new graphs.
print(my_relu(tf.constant(-2.5))) # Signature matches `tf.constant(5.5)`.
print(my_relu(tf.constant([-1., 1.]))) # Signature matches `tf.constant([3., -3.])`.
tf.Tensor(0.0, shape=(), dtype=float32)
tf.Tensor([0. 1.], shape=(2,), dtype=float32)

Porque ele é apoiado por vários gráficos, a Function é polimórfica. Que lhe permite suportar tipos mais entrada do que uma única tf.Graph poderiam representar, bem como para otimizar cada tf.Graph para um melhor desempenho.

# There are three `ConcreteFunction`s (one for each graph) in `my_relu`.
# The `ConcreteFunction` also knows the return type and shape!
print(my_relu.pretty_printed_concrete_signatures())
my_relu(x)
  Args:
    x: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()

my_relu(x=[1, -1])
  Returns:
    float32 Tensor, shape=(2,)

my_relu(x)
  Args:
    x: float32 Tensor, shape=(2,)
  Returns:
    float32 Tensor, shape=(2,)

usando tf.function

Até agora, você aprendeu como converter uma função Python em um gráfico simplesmente usando tf.function como um decorador ou wrapper. Mas, na prática, ficando tf.function ao trabalho corretamente pode ser complicado! Nas seções seguintes, você vai aprender como você pode fazer seu trabalho de código conforme esperado com tf.function .

Execução de gráfico vs. execução ansiosa

O código em uma Function pode ser executado tanto ansiosa e como um gráfico. Por padrão, Function executa seu código como um gráfico:

@tf.function
def get_MSE(y_true, y_pred):
  sq_diff = tf.pow(y_true - y_pred, 2)
  return tf.reduce_mean(sq_diff)
y_true = tf.random.uniform([5], maxval=10, dtype=tf.int32)
y_pred = tf.random.uniform([5], maxval=10, dtype=tf.int32)
print(y_true)
print(y_pred)
tf.Tensor([6 4 3 1 5], shape=(5,), dtype=int32)
tf.Tensor([6 7 9 9 1], shape=(5,), dtype=int32)
get_MSE(y_true, y_pred)
<tf.Tensor: shape=(), dtype=int32, numpy=25>

Para verificar se a Function do gráfico 's está fazendo o mesmo cálculo como função Python equivalente, você pode fazê-lo executar ansiosamente com tf.config.run_functions_eagerly(True) . Este é um interruptor que desliga Function capacidade de criar e gráficos executados, em vez de executar o código normalmente.

tf.config.run_functions_eagerly(True)
get_MSE(y_true, y_pred)
<tf.Tensor: shape=(), dtype=int32, numpy=25>
# Don't forget to set it back when you are done.
tf.config.run_functions_eagerly(False)

No entanto, Function pode se comportar de forma diferente sob gráfico e execução ansioso. O Python print função é um exemplo de como estes dois modos diferentes. Vamos verificar o que acontece quando você inserir um print declaração para a sua função e chamá-lo repetidamente.

@tf.function
def get_MSE(y_true, y_pred):
  print("Calculating MSE!")
  sq_diff = tf.pow(y_true - y_pred, 2)
  return tf.reduce_mean(sq_diff)

Observe o que está impresso:

error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
Calculating MSE!

O resultado é surpreendente? get_MSE única impresso uma vez apesar de ter sido chamado três vezes.

Para explicar, a print instrução é executada quando Function executa o código de origem para criar o gráfico em um processo conhecido como "rastreamento" . Rastreamento captura as operações TensorFlow em um gráfico e print não é capturado no gráfico. Esse gráfico é então executado para todas as três chamadas sem nunca executar o código Python novamente.

Para verificar a integridade, vamos desligar a execução do gráfico para comparar:

# Now, globally set everything to run eagerly to force eager execution.
tf.config.run_functions_eagerly(True)
# Observe what is printed below.
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
Calculating MSE!
Calculating MSE!
Calculating MSE!
tf.config.run_functions_eagerly(False)

print é um efeito colateral Python, e há outras diferenças que você deve estar ciente de quando a conversão de uma função em uma Function .

Execução não estrita

A execução do gráfico executa apenas as operações necessárias para produzir os efeitos observáveis, que incluem:

  • O valor de retorno da função
  • Efeitos colaterais bem conhecidos documentados, como:
    • Operações de entrada / saída, como tf.print
    • Operações de depuração, tais como as funções de declaração no tf.debugging
    • Mutações de tf.Variable

Esse comportamento é geralmente conhecido como "execução não estrita" e difere da execução antecipada, que percorre todas as operações do programa, necessárias ou não.

Em particular, a verificação de erros de tempo de execução não conta como um efeito observável. Se uma operação for ignorada por ser desnecessária, ela não poderá gerar nenhum erro de tempo de execução.

No exemplo a seguir, a operação "desnecessário" tf.gather é ignorada durante a execução do gráfico, assim que o erro de execução InvalidArgumentError não é gerado como seria na execução ansioso. Não confie no surgimento de um erro ao executar um gráfico.

def unused_return_eager(x):
  # Get index 1 will fail when `len(x) == 1`
  tf.gather(x, [1]) # unused 
  return x

try:
  print(unused_return_eager(tf.constant([0.0])))
except tf.errors.InvalidArgumentError as e:
  # All operations are run during eager execution so an error is raised.
  print(f'{type(e).__name__}: {e}')
tf.Tensor([0.], shape=(1,), dtype=float32)
@tf.function
def unused_return_graph(x):
  tf.gather(x, [1]) # unused
  return x

# Only needed operations are run during graph exection. The error is not raised.
print(unused_return_graph(tf.constant([0.0])))
tf.Tensor([0.], shape=(1,), dtype=float32)

tf.function melhores práticas

Pode levar algum tempo para se acostumar com o comportamento da Function . Para começar rapidamente, os usuários de primeira viagem deve brincar com a decoração funções brinquedo com @tf.function para obter experiência com indo de ansioso para execução gráfico.

Projetando para tf.function pode ser sua melhor aposta para escrever programas TensorFlow gráfico compatível. Aqui estão algumas dicas:

Vendo a aceleração

tf.function geralmente melhora o desempenho do seu código, mas a quantidade de aceleração depende do tipo de computação que você executa. Pequenos cálculos podem ser dominados pela sobrecarga de chamar um gráfico. Você pode medir a diferença de desempenho da seguinte forma:

x = tf.random.uniform(shape=[10, 10], minval=-1, maxval=2, dtype=tf.dtypes.int32)

def power(x, y):
  result = tf.eye(10, dtype=tf.dtypes.int32)
  for _ in range(y):
    result = tf.matmul(x, result)
  return result
print("Eager execution:", timeit.timeit(lambda: power(x, 100), number=1000))
Eager execution: 2.286959676999686
power_as_graph = tf.function(power)
print("Graph execution:", timeit.timeit(lambda: power_as_graph(x, 100), number=1000))
Graph execution: 0.6817171659999985

tf.function é comumente usado para acelerar ciclos de treinamento, e você pode aprender mais sobre ele em escrever um loop de formação a partir do zero com Keras.

Desempenho e trade-offs

Os gráficos podem acelerar seu código, mas o processo de criá-los tem alguma sobrecarga. Para algumas funções, a criação do gráfico leva mais tempo do que a execução do gráfico. Esse investimento geralmente é reembolsado rapidamente com o aumento de desempenho das execuções subsequentes, mas é importante estar ciente de que as primeiras etapas de qualquer treinamento de modelo grande podem ser mais lentas devido ao rastreamento.

Não importa o tamanho do seu modelo, você deseja evitar rastreios com frequência. Os tf.function guia discute como especificações conjunto de entrada e argumentos utilização de tensores para evitar refazer. Se você descobrir que está obtendo um desempenho anormalmente ruim, é uma boa ideia verificar se está refazendo acidentalmente.

Quando é uma Function de rastreamento?

Para descobrir quando a sua Function está traçando, adicione um print declaração ao seu código. Como regra geral, Function irá executar a print afirmação cada vez que traça.

@tf.function
def a_function_with_python_side_effect(x):
  print("Tracing!") # An eager-only side effect.
  return x * x + tf.constant(2)

# This is traced the first time.
print(a_function_with_python_side_effect(tf.constant(2)))
# The second time through, you won't see the side effect.
print(a_function_with_python_side_effect(tf.constant(3)))
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
tf.Tensor(11, shape=(), dtype=int32)
# This retraces each time the Python argument changes,
# as a Python argument could be an epoch count or other
# hyperparameter.
print(a_function_with_python_side_effect(2))
print(a_function_with_python_side_effect(3))
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
Tracing!
tf.Tensor(11, shape=(), dtype=int32)

Novos argumentos Python sempre disparam a criação de um novo gráfico, daí o rastreamento extra.

Próximos passos

Você pode aprender mais sobre tf.function na página de referência API e seguindo o desempenho melhor com tf.function guia.