![]() | ![]() | ![]() | ![]() |
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 umtf.Graph
. - Uma
Function
gerencia um cache deConcreteFunction
escolhe o certo para suas entradas. -
tf.function
envolve uma função Python, retornando um objetoFunction
. - O rastreamento cria um
tf.Graph
e o envolve em umConcreteFunction
, 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 eattr
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
emtf.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á queFunction
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
outraining=True
ornonlinearity='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
chamadastf.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
: sey
for um Tensor, converta paratf.while_loop
. No caso especial em quey
é umtf.data.Dataset
, uma combinação de operaçõestf.data.Dataset
é gerada. -
while <condition>
: se<condition>
for um Tensor, converta paratf.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 .