function.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Melhor desempenho com tf.function

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

No TensorFlow 2, a execução eager está ativada por padrão. A interface do usuário é intuitiva e flexível (executar 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 criar grafos de seus programas. É uma ferramenta de transformação que cria grafos de dataflow independentes do Python a partir do seu código Python. Isto vai ajudar você a criar modelos portáteis e de alto desempenho, além de ser necessário para poder usar o SavedModel.

Este guia vai ajudar você a conceituar como tf.function funciona nos bastidores, para que você possa usá-lo de maneira eficaz.

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

• Depure no modo eager e decore com @tf.function.

• Não conte com os efeitos colaterais do Python, como mutação de objetos ou appends de listas.

• tf.function funciona melhor com as ops do TensorFlow. As chamadas NumPy e Python são convertidas em constantes.

Configuração

import tensorflow as tf

Defina uma função helper para demonstrar os tipos de erros que você poderá 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 que você define (por exemplo, aplicando o decorador @tf.function) é como uma operação do TensorFlow core: você pode executá-la de forma eager; você pode computar 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.]]

v = tf.Variable(1.0)
with tf.GradientTape() as tape:
  result = add(v, 1.0)
tape.gradient(result, v)

Você pode usar uma Function dentro de outra Function.

@tf.function
def dense_layer(x, w, b):
  return add(tf.matmul(x, w), b)

dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))

Uma Function pode ser mais rápida que o código eager, especialmente em grafos com muitas pequenas ops. Mas para grafos com algumas poucas operações caras (como convoluções), talvez você não perceba muita diferença.

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

Tracing

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

O que é "tracing"?

Uma Function executa seu programa num TensorFlow Graph. No entanto, um tf.Graph não pode ser usado para representar tudo o que você poderia escrever num programa TensorFlow em modo eager. Por exemplo, Python suporta polimorfismo, mas tf.Graph exige que suas entradas tenham um tipo de dados e dimensão especificados. Ou você talvez queira realizar tarefas paralelas, como ler argumentos de linha de comando, gerar um erro ou trabalhar com um objeto Python mais complexo; nenhuma dessas coisas irá rodar num tf.Graph.

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

1. No primeiro estágio, chamado de "tracing", 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. No segundo estágio é executado um tf.Graph que contém tudo o que foi adiado na primeira etapa. Este estágio é muito mais rápido que o estágio de tracing.

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

Quando Function decide fazer tracing, o estágio de tracing é imediatamente seguido pelo segundo estágio, portanto, chamar Function cria e executa o tf.Graph. Mais tarde, você verá como executar apenas o estágio de tracing com get_concrete_function.

Quando você passa argumentos de tipos diferentes para um Function, ambos os 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()

Observe que se você chamar repetidamente uma Function com o mesmo tipo de argumento, o TensorFlow ignorará o estágio de tracing e reutilizará um grafo que tenha passado pelo tracing anteriormente, pois o grafo gerado seria idêntico.

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

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

print(double.pretty_printed_concrete_signatures())

Até aqui, você viu que tf.function cria uma camada de despacho dinâmico e em cache sobre a lógica de tracing de grafos do TensorFlow. Para sermos mais específicos sobre a terminologia:

• Um tf.Graph é a representação bruta, independente de linguagem e portátil de uma computação do TensorFlow.

• Uma ConcreteFunction é um wrapper que envolve um tf.Graph.

• Uma Function gerencia um cache de ConcreteFunction e escolhe a função correta para suas entradas.

• tf.function é um wrapper que envolve uma função Python, retornando um objeto Function.

• Tracing cria um tf.Graph e o empacota num ConcreteFunction , também conhecido como trace.

Regras do tracing

Quando chamada, uma Function combina os argumentos de chamada com as ConcreteFunction existentes usando tf.types.experimental.TraceType de cada argumento. Se uma ConcreteFunction correspondente for encontrada, a chamada será enviada para ela. Se nenhuma correspondência for encontrada, será feito o tracing de uma nova ConcreteFunction.

Se forem encontradas múltiplas correspondências, a assinatura mais específica será escolhida. A correspondência é feita por meio de subtipos, de forma similar a chamadas de função comuns em C++ ou Java. Por exemplo, TensorShape([1, 2]) é um subtipo de TensorShape([None, None]) e, portanto, uma chamada para tf.function com TensorShape([1, 2]) pode ser despachada para ConcreteFunction produzida com TensorShape([None, None]) mas se um ConcreteFunction com TensorShape([1, None]) também existir, ele será priorizado por ser mais específico.

O TraceType é determinado a partir de argumentos de entrada da seguinte forma:

• Para Tensor, o tipo é parametrizado pelo dtype e shape do Tensor; os formatos classificados são um subtipo de formatos não classificados; dimensões fixas são um subtipo de dimensões desconhecidas

• Para Variable, o tipo é semelhante a Tensor, mas também inclui um ID de recurso exclusivo da variável, necessário para interligar corretamente as dependências de controle

• Para valores primitivos do Python, o tipo corresponde ao próprio valor. Por exemplo, o TraceType do valor 3 é LiteralTraceType<3>, não int.

• Para containers ordenados em Python, como list e tuple, etc., o tipo é parametrizado pelos tipos de seus elementos; por exemplo, o tipo de [1, 2] é ListTraceType<LiteralTraceType<1>, LiteralTraceType<2>> e o tipo de [2, 1] é ListTraceType<LiteralTraceType<2>, LiteralTraceType<1>> que é diferente.

• Para mapeamentos Python como dict, o tipo também é um mapeamento das mesmas chaves, mas para os tipos de valores, em vez dos valores reais. Por exemplo, o tipo de {1: 2, 3: 4}, é MappingTraceType<<KeyValue<1, LiteralTraceType<2>>>, <KeyValue<3, LiteralTraceType<4>>>>. No entanto, diferentemente dos containers ordenados, {1: 2, 3: 4} e {3: 4, 1: 2} possuem tipos equivalentes.

• Para objetos Python que implementam o método __tf_tracing_type__, o tipo é o que o método retorna

• Para quaisquer outros objetos Python, o tipo é um TraceType genérico, seu procedimento correspondente é:

  • Primeiro ele verifica se o objeto é o mesmo usado no tracing anterior (usando id() ou is, do Python). Observe que isto ainda vai corresponder se o objeto tiver sido alterado; portanto, se você usar objetos Python como argumentos tf.function, é melhor que sejam objetos imutáveis.

  • Em seguida, verifica se o objeto é igual ao objeto usado no tracing anterior (usando == do Python).

  Observe que este procedimento apenas mantém uma referência fraca para o objeto e, portanto, só funciona enquanto o objeto estiver no escopo/não for excluído.)

Observação: TraceType é baseado nos parâmetros de entrada de Function, portanto, alterações apenas nas variáveis ​​globais e livres não criarão um novo trace. Consulte esta seção para conhecer as práticas recomendadas ao lidar com variáveis ​​globais e livres do Python.

Controle do retracing

O retracing, que ocorre quando sua Function cria mais de um trace, ajuda a garantir que o TensorFlow gere grafos corretos para cada conjunto de entradas. No entanto, o tracing é uma operação cara! Se sua Function fizer retracing de um novo gráfico para cada chamada, você vai perceber que seu código roda mais lentamente do que se você não usasse tf.function.

Para controlar o comportamento do tracing, você pode usar as seguintes técnicas:

Passe uma input_signature fixa para tf.function

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)

print(next_collatz(tf.constant([1, 2])))
# You specified a 1-D tensor in the input signature, so this should fail.
with assert_raises(TypeError):
  next_collatz(tf.constant([[1, 2], [3, 4]]))

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

Use dimensões desconhecidas para maior flexibilidade

Como o TensorFlow compara tensores com base no seu formato, usar uma dimensão None como curinga permitirá que as Function reutilizem traces para entradas de tamanhos variáveis. Entradas de tamanhos variáveis ​​podem ocorrer se você tiver sequências de comprimentos diferentes ou imagens de tamanhos diferentes para cada lote (veja os tutoriais 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])))

Passe tensores em vez de literais python

Freqüentemente, argumentos do Python são usados ​​para controlar hiperparâmetros e construções de grafos, por exemplo, num_layers=10 ou training=True ou nonlinearity='relu'. Portanto, se o argumento Python mudar, faz sentido que você tenha que fazer o retracing do grafo.

No entanto, é possível que um argumento Python não esteja sendo usado para controlar a construção do grafo. Nesses casos, uma alteração no valor Python pode desencadear um retracing desnecessário. Veja, por exemplo, este loop de treinamento, que o AutoGraph irá desenrolar dinamicamente. Apesar dos múltiplos traces, o gráfico gerado é na verdade idêntico, portanto, fazer retracing aqui é 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))

Se você precisar forçar o retracing, crie uma nova Function. É garantido que objetos Function separados não compartilhem traces.

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

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

Use o protocolo de tracing

Sempre que possível, você deve dar preferência a converter o tipo Python em tf.experimental.ExtensionType. Além disso, o TraceType de um ExtensionType é o tf.TypeSpec associado a ele. Portanto, se necessário, você pode simplesmente sobrepor o tf.TypeSpec padrão para assumir o controle de um Tracing Protocol do ExtensionType. Consulte a seção Personalizando o TypeSpec do ExtensionType no guia Extension types para mais detalhes.

Caso contrário, para ter controle direto sobre quando Function deve fazer retracing quanto a um tipo específico do Python, você mesmo pode implementar o Tracing Protocol para ela.

@tf.function
def get_mixed_flavor(fruit_a, fruit_b):
  return fruit_a.flavor + fruit_b.flavor

class Fruit:
  flavor = tf.constant([0, 0])

class Apple(Fruit):
  flavor = tf.constant([1, 2])

class Mango(Fruit):
  flavor = tf.constant([3, 4])

# As described in the above rules, a generic TraceType for `Apple` and `Mango`
# is generated (and a corresponding ConcreteFunction is traced) but it fails to
# match the second function call since the first pair of Apple() and Mango()
# have gone out out of scope by then and deleted.
get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function
get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function again

# However, each subclass of the `Fruit` class has a fixed flavor, and you
# can reuse an existing traced concrete function if it was the same
# subclass. Avoiding such unnecessary tracing of concrete functions
# can have significant performance benefits.

class FruitTraceType(tf.types.experimental.TraceType):
  def __init__(self, fruit):
    self.fruit_type = type(fruit)
    self.fruit_value = fruit

  def is_subtype_of(self, other):
      # True if self subtypes `other` and `other`'s type matches FruitTraceType.
      return (type(other) is FruitTraceType and
              self.fruit_type is other.fruit_type)

  def most_specific_common_supertype(self, others):
      # `self` is the specific common supertype if all input types match it.
      return self if all(self == other for other in others) else None

  def placeholder_value(self, placeholder_context=None):
      # Use the fruit itself instead of the type for correct tracing.
      return self.fruit_value

  def __eq__(self, other):
    return type(other) is FruitTraceType and self.fruit_type == other.fruit_type

  def __hash__(self):
    return hash(self.fruit_type)

class FruitWithTraceType:

  def __tf_tracing_type__(self, context):
    return FruitTraceType(self)

class AppleWithTraceType(FruitWithTraceType):
  flavor = tf.constant([1, 2])

class MangoWithTraceType(FruitWithTraceType):
  flavor = tf.constant([3, 4])

# Now if you try calling it again:
get_mixed_flavor(AppleWithTraceType(), MangoWithTraceType()) # Traces a new concrete function
get_mixed_flavor(AppleWithTraceType(), MangoWithTraceType()) # Re-uses the traced concrete function

Obtendo funções concretas

Cada vez que o tracing de uma função é realizado, 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")))

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

Imprimir uma ConcreteFunction mostra um resumo de seus argumentos de entrada (com tipos) e seu tipo de saída.

print(double_strings)

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

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

Usar um trace concreto com tipos incompatíveis gerará um erro

with assert_raises(tf.errors.InvalidArgumentError):
  double_strings(tf.constant(1))

Você talvez perceba que os argumentos Python recebem tratamento especial na assinatura de entrada de uma função concreta. Antes do TensorFlow 2.3, os argumentos Python eram simplesmente removidos da assinatura da função concreta. A partir do TensorFlow 2.3, os argumentos Python permanecem na assinatura, mas são limitados a aceitar o valor definido durante o tracing.

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

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

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

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

Obtendo grafos

Cada função concreta é um wrapper em torno de um tf.Graph que pode ser chamado. 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}')

Depuração

Em geral, depurar código é mais fácil no modo eager do que dentro de tf.function. Você deve garantir que seu código seja executado sem erros no modo eager antes de decorar com tf.function. Para auxiliar 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 aparecem apenas em tf.function, aqui estão algumas dicas:

• Chamadas print do Python são executadas apenas durante o tracing, ajudando você a rastrear quando ocorrer o (re)tracing da sua função.

• As chamadas tf.print serão executadas sempre 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 (o depurador Python ) pode ajudá-lo a entender o que está acontecendo durante o tracing. (Ressalva: o pdb vai mandar você para o código-fonte transformado pelo AutoGraph.)

Transformações AutoGraph

AutoGraph é uma biblioteca ativada por padrão em tf.function e transforma um subconjunto de código eager do Python em ops do TensorFlow compatíveis com grafos. Isto inclui operações de controle de fluxo como if, for, while.

Ops do TensorFlow como tf.cond e tf.while_loop continuam funcionando, mas o controle de fluxo geralmente é 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]))

Se você tiver curiosidade, pode inspecionar o código gerado pelo AutoGraph.

print(tf.autograph.to_code(f.python_function))

Condicionais

O AutoGraph converterá algumas instruções if <condition> em chamadas tf.cond equivalentes. Esta substituição é feita se <condition> for um Tensor. Caso contrário, a instrução if será executada como uma condicional Python.

Uma expressão condicional Python é executada durante o tracing, portanto, exatamente um ramo da expressão condicional será adicionada ao grafo. Sem o AutoGraph, este grafo rastreado não seria capaz de seguir o ramo alternativo se houvesse controle de fluxo dependente de dados.

tf.cond faz o tracing e adiciona ambos os ramos da expressão condicional ao grafo, selecionando dinamicamente um ramo em tempo de execução. O tracing pode ter efeitos colaterais indesejados; veja Efeitos de tracing do AutoGraph 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))

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

Loops

O AutoGraph converterá algumas instruções for e while em operações de loop equivalentes do TensorFlow, como tf.while_loop. Se não for convertido, o loop for ou while será executado como um loop 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 é um tf.data.Dataset, uma combinação de operações de tf.data.Dataset é gerada.

• while <condition>: se <condition> for um Tensor, converta para tf.while_loop.

Um loop Python é executado durante o tracing, adicionando ops adicionais ao tf.Graph para cada iteração do loop.

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

Veja a documentação de referência para conhecer restrições adicionais nas instruções for e while convertidas pelo AutoGraph.

Loop sobre dados Python

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

Se você quiser encapsular todo o loop de treinamento em tf.function, a maneira mais segura de fazer isso é encapsular seus dados usando um tf.data.Dataset como wrapper 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)))

Ao encapsular dados Python/NumPy num dataset, lembre-se de tf.data.Dataset.from_generator versus tf.data.Dataset.from_tensor_slices. O primeiro manterá os dados em Python e os buscará via tf.py_function, o que pode ter implicações no desempenho, enquanto o último agrupará uma cópia dos dados como uma grande tf.constant() no grafo, o que pode ter implicações no uso de memória.

Ler 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 a necessidade de envolver o Python. Para saber mais, consulte o guia tf.data: Criando pipelines de entrada no TensorFlow.

Acumulando valores em um loop

Um padrão comum é acumular valores intermediários de um loop. Normalmente, isso é feito anexando-se 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 num 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]))

Limitações

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

Executando efeitos colaterais do Python

Os efeitos colaterais, como impressão, anexação a listas e mutação de variáveis globais, podem se comportar inesperadamente dentro de uma Function, às vezes executando duas vezes ou nunca. Isto só acontece na primeira vez que você chama uma Function com um conjunto de entradas. Posteriormente, o tf.Graph traçado é reexecutado, e não 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 traces. Caso contrário, 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 será executado pelo runtime 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)

Se você quiser executar o código Python durante cada chamada de uma Function, tf.py_function é uma saída. A desvantagem de tf.py_function é que ela não é portátil nem tem bom desempenho, não pode ser salva com SavedModel e não funciona bem em configurações distribuídas (multi-GPU, TPU). Além disso, como tf.py_function precisa ser interligado ao grafo, ela converte todas as entradas/saídas em tensores.

Alterando variáveis ​​globais e livres do Python

Alterar as variáveis ​​globais e livres do Python conta como um efeito colateral do Python, então só acontece durante o tracing.

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

Às vezes, comportamentos inesperados são muito difíceis de perceber. No exemplo abaixo, o counter tem como objetivo salvaguardar o incremento de uma variável. No entanto, como é um inteiro Python e não um objeto TensorFlow, seu valor é capturado durante o primeiro trace. Quando tf.function for usado, assign_add será registrado incondicionalmente no grafo subjacente. Portanto v aumentará em 1, todas as vezes que a tf.function for chamada. Esse problema é comum entre usuários que tentam migrar seu código Tensorflow no modo grafo para o Tensorflow 2 usando decoradores tf.function, quando os efeitos colaterais do Python (o counter no exemplo) são usados ​​para determinar quais operações executar (assign_add no exemplo). Normalmente, os usuários percebem isso somente depois de verem resultados numéricos suspeitos ou um desempenho significativamente inferior ao esperado (por exemplo, se a operação protegida for muito ineficiente).

class Model(tf.Module):
  def __init__(self):
    self.v = tf.Variable(0)
    self.counter = 0

  @tf.function
  def __call__(self):
    if self.counter == 0:
      # A python side-effect
      self.counter += 1
      self.v.assign_add(1)

    return self.v

m = Model()
for n in range(3):
  print(m().numpy()) # prints 1, 2, 3

Uma solução alternativa para alcançar o comportamento esperado é usar tf.init_scope para transferir as operações para fora do grafo de função. Isto garante que o incremento da variável seja feito apenas uma vez durante o tempo de tracing. Deve-se observar que init_scope tem outros efeitos colaterais, incluindo fluxo de controle limpo e fita gradiente. Às vezes, o uso de init_scope pode se tornar muito complexo para ser gerenciado de forma realista.

class Model(tf.Module):
  def __init__(self):
    self.v = tf.Variable(0)
    self.counter = 0

  @tf.function
  def __call__(self):
    if self.counter == 0:
      # Lifts ops out of function-building graphs
      with tf.init_scope():
        self.counter += 1
        self.v.assign_add(1)

    return self.v

m = Model()
for n in range(3):
  print(m().numpy()) # prints 1, 1, 1

Em suma, como regra geral, você deve evitar a mutação de objetos python, como números inteiros ou containers, como listas que residem fora de Function. Em vez disso, use argumentos e objetos do TF. Por exemplo, a seção "Acumulando valores em um loop" traz um exemplo de como 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 a mesma ConcreteFunction.

Usando iteradores e geradores Python

Muitos recursos do Python, como geradores e iteradores, dependem do runtime do Python para controlar o estado. Em geral, embora esses construtos funcionem conforme o esperado no modo eager, elas são exemplos de efeitos colaterais do Python e, portanto, só acontecem durante o tracing.

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

Assim como o TensorFlow possui um tf.TensorArray especializado para construtos de lista, ele possui um tf.data.Iterator especializado para construtos de iteração. Veja a seção sobre transformações do 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)

Todas as saídas de uma tf.function devem ser valores de retorno

Com exceção de tf.Variable, uma tf.function deve retornar todas as suas saídas. Tentar acessar diretamente qualquer tensor de uma função sem passar pelos valores de retorno causa "vazamentos".

Por exemplo, a função abaixo "vaza" o tensor a através da variável global x declarada em Python:

x = None

@tf.function
def leaky_function(a):
  global x
  x = a + 1  # Bad - leaks local tensor
  return a + 2

correct_a = leaky_function(tf.constant(1))

print(correct_a.numpy())  # Good - value obtained from function's returns
try:
  x.numpy()  # Bad - tensor leaked from inside the function, cannot be used here
except AttributeError as expected:
  print(expected)

Isso é verdade mesmo que o valor vazado também seja retornado:

@tf.function
def leaky_function(a):
  global x
  x = a + 1  # Bad - leaks local tensor
  return x  # Good - uses local tensor

correct_a = leaky_function(tf.constant(1))

print(correct_a.numpy())  # Good - value obtained from function's returns
try:
  x.numpy()  # Bad - tensor leaked from inside the function, cannot be used here
except AttributeError as expected:
  print(expected)

@tf.function
def captures_leaked_tensor(b):
  b += x  # Bad - `x` is leaked from `leaky_function`
  return b

with assert_raises(TypeError):
  captures_leaked_tensor(tf.constant(2))

Geralmente, vazamentos como esses ocorrem quando você usa instruções ou estruturas de dados do Python. Além de vazar tensores inacessíveis, tais instruções provavelmente também estão erradas porque contam como efeitos colaterais do Python e não têm garantia de execução a cada chamada de função.

Maneiras comuns de vazar tensores locais também incluem a mutação de uma coleção Python externa ou de um objeto:

class MyClass:

  def __init__(self):
    self.field = None

external_list = []
external_object = MyClass()

def leaky_function():
  a = tf.constant(1)
  external_list.append(a)  # Bad - leaks tensor
  external_object.field = a  # Bad - leaks tensor

Funções tf.function recursivas não são suportadas

Function recursivas não são suportadas e podem causar loops infinitos. Por exemplo,

@tf.function
def recursive_fn(n):
  if n > 0:
    return recursive_fn(n - 1)
  else:
    return 1

with assert_raises(Exception):
  recursive_fn(tf.constant(5))  # Bad - maximum recursion error.

Mesmo que uma Function recursiva pareça funcionar, ocorrerá tracing da função Python várias vezes e isto poderá ter implicações no desempenho. Por exemplo,

@tf.function
def recursive_fn(n):
  if n > 0:
    print('tracing')
    return recursive_fn(n - 1)
  else:
    return 1

recursive_fn(5)  # Warning - multiple tracings

Problemas conhecidos

Se a sua Function não estiver sendo avaliada corretamente, o erro pode ser explicado por esses problemas conhecidos que devem ser corrigidos no futuro.

Depender de variáveis ​​globais e livres do Python

Function cria uma nova ConcreteFunction quando chamada com um novo valor de um argumento Python. No entanto, ele não faz isso para o closure do Python, nem valores globais ou não locais dessa Function. Se o valor deles mudar entre as chamadas para Function, a Function ainda usará os valores que tinham quando ela passou pelo tracing. Isso difere da forma como funcionam as funções comuns do Python.

Por esse motivo, você deve seguir um estilo de programação funcional que use argumentos em vez de closures 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))

print("Updating the value of `foo` to 100!")
foo = 100
print("Buggy:", buggy_add())  # Did not change!
print("Correct:", recommended_add(foo))

Outra maneira de atualizar um valor global é torná-lo uma tf.Variable e usar o método Variable.assign.

@tf.function
def variable_add():
  return 1 + foo

foo = tf.Variable(1)
print("Variable:", variable_add())

print("Updating the value of `foo` to 100!")
foo.assign(100)
print("Variable:", variable_add())

Depender de objetos Python

Você pode passar objetos Python personalizados como argumentos para tf.function mas isso traz algumas limitações.

Para cobertura máxima de recursos, considere transformar os objetos em tipos de extensão antes de passá-los para tf.function. Você também pode usar primitivos Python e estruturas compatíveis com tf.nest.

No entanto, conforme abordado nas regras de tracing, quando um TraceType personalizado não é fornecido pela classe Python personalizada, a tf.function é obrigada a usar igualdade baseada em instância, o que significa que não criará um novo trace quando você passar o mesmo objeto com atributos modificados.

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

print("Adding bias!")
simple_model.bias += 5.0
print(evaluate(simple_model, x))  # Didn't change :(

Usar a mesma Function para avaliar a instância modificada do modelo será problemático, pois ele ainda possui o mesmo TraceType baseado em instância do modelo original.

Por esse motivo, é recomendável escrever sua Function para evitar depender de atributos de objetos mutáveis ​​ou implementar o Protocolo de Tracing para os objetos para informar Function sobre tais atributos.

Se isso não for possível, uma solução alternativa é criar novos objetos Function cada vez que você modificar seu objeto para forçar o 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))

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`.

Como o retracing pode ter um custo elevado, você pode usar objetos tf.Variable como atributos de objeto, que podem ser alterados (mas não trocados, cuidado!) para um efeito semelhante sem a necessidade 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))

print("Adding bias!")
better_model.bias.assign_add(5.0)  # Note: instead of better_model.bias += 5
print(evaluate(better_model, x))  # This works!

Criação de tf.Variables

Function suporta apenas objetos tf.Variable singleton, criados uma vez na primeira chamada e reutilizados nas chamadas de função subsequentes. O trecho de código abaixo criaria um novo tf.Variable em cada chamada de função, o que resulta numa exceção ValueError.

Exemplo:

@tf.function
def f(x):
  v = tf.Variable(1.0)
  return v

with assert_raises(ValueError):
  f(1.0)

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

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

Uso com múltiplos otimizadores Keras

Você talvez encontre o erro ValueError: tf.function only supports singleton tf.Variables created on the first call. ao usar mais de um otimizador Keras com um tf.function. Esse erro ocorre porque os otimizadores criam tf.Variables internamente 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)

Se você precisar trocar o otimizador durante o treinamento, uma solução alternativa é criar uma nova Function para cada otimizador, chamando 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)
train_step_2 = tf.function(train_step)
for i in range(10):
  if i % 2 == 0:
    train_step_1(w, x, y, opt1)
  else:
    train_step_2(w, x, y, opt2)

Uso com múltiplos modelos Keras

Você talvez também encontre ValueError: tf.function only supports singleton tf.Variables created on the first call. ao passar diferentes instâncias de modelo para a mesma Function.

Este erro ocorre porque os modelos Keras (que não têm seu formato de entrada definido) e as camadas Keras criam objetos tf.Variables quando são chamados 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.

Leituras adicionais

Para saber como exportar e carregar uma Function, consulte o guia do SavedModel. Para saber mais sobre otimizações de grafos executadas após o tracing, consulte o guia do Grappler. Para saber como otimizar seu pipeline de dados e fazer um profiling do seu modelo, consulte o Guia do profiler.