¡Google I / O regresa del 18 al 20 de mayo! Reserva espacio y crea tu horario Regístrate ahora

Mejor rendimiento con tf.function

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

En TensorFlow 2, la ejecución ansiosa está activada de forma predeterminada. La interfaz de usuario es intuitiva y flexible (ejecutar operaciones puntuales es mucho más fácil y rápido), pero esto puede producirse a expensas del rendimiento y la capacidad de implementación.

Puede usar tf.function para hacer gráficos a partir de sus programas. Es una herramienta de transformación que crea gráficos de flujo de datos independientes de Python a partir de su código Python. Esto le ayudará a crear modelos portátiles y de rendimiento, y es necesario utilizar SavedModel .

Esta guía lo ayudará a conceptualizar cómo funciona tf.function debajo del capó para que pueda usarlo de manera efectiva.

Las principales conclusiones y recomendaciones son:

  • Depura en modo ansioso, luego decora con @tf.function .
  • No confíe en los efectos secundarios de Python, como la mutación de objetos o la lista de anexos.
  • tf.function funciona mejor con las operaciones de TensorFlow; Las llamadas a NumPy y Python se convierten en constantes.

Configuración

import tensorflow as tf

Defina una función auxiliar para demostrar los tipos de errores que puede 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))

Lo esencial

Uso

Una Function que defina (por ejemplo, aplicando el decorador @tf.function ) es como una operación central de TensorFlow: puede ejecutarla con entusiasmo; puedes calcular gradientes; y así.

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

Puede utilizar Function s dentro de otras 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 pueden ser más rápidas que el código ávido, especialmente para gráficos con muchas operaciones pequeñas. Pero para gráficos con algunas operaciones costosas (como convoluciones), es posible que no vea mucha aceleración.

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

Rastreo

Esta sección expone cómo Function la Function bajo el capó, incluidos los detalles de implementación que pueden cambiar en el futuro . Sin embargo, una vez que comprenda por qué y cuándo ocurre el rastreo, ¡es mucho más fácil usar tf.function efectiva!

¿Qué es "rastreo"?

Una Function ejecuta su programa en un gráfico de TensorFlow . Sin embargo, un tf.Graph no puede representar todas las cosas que escribirías en un programa TensorFlow ansioso. Por ejemplo, Python admite polimorfismo, pero tf.Graph requiere que sus entradas tengan un tipo de datos y una dimensión específicos. O puede realizar tareas secundarias como leer argumentos en la línea de comandos, generar un error o trabajar con un objeto Python más complejo; ninguna de estas cosas puede ejecutarse en un tf.Graph .

Function cierra esta brecha al separar su código en dos etapas:

1) En la primera etapa, denominada " rastreo ", la Function crea un nuevo tf.Graph . El código de Python se ejecuta normalmente, pero todas las operaciones de TensorFlow (como agregar dos tensores) se aplazan : son capturadas por tf.Graph y no se ejecutan.

2) En la segunda etapa, se tf.Graph un tf.Graph que contiene todo lo que fue aplazado en la primera etapa. Esta etapa es mucho más rápida que la etapa de rastreo.

Dependiendo de sus entradas, Function no siempre ejecutará la primera etapa cuando se la llame. Consulte "Reglas de rastreo" a continuación para tener una mejor idea de cómo toma esa determinación. Saltarse la primera etapa y solo ejecutar la segunda etapa es lo que le brinda el alto rendimiento de TensorFlow.

Cuando la Function decide rastrear, la etapa de rastreo es seguida inmediatamente por la segunda etapa, por lo que llamar a la Function crea y ejecuta tf.Graph . Más adelante verá cómo puede ejecutar solo la etapa de seguimiento con get_concrete_function .

Cuando pasamos argumentos de diferentes tipos a una Function , se ejecutan ambas etapas:

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

Tenga en cuenta que si llama repetidamente a una Function con el mismo tipo de argumento, TensorFlow omitirá la etapa de seguimiento y reutilizará un gráfico previamente trazado, ya que el gráfico generado sería idéntico.

# This doesn't print 'Tracing with ...'
print(double(tf.constant("b")))
tf.Tensor(b'bb', shape=(), dtype=string)

Puede usar pretty_printed_concrete_signatures() para ver todos los rastros disponibles:

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

Hasta ahora, has visto que tf.function crea una capa de distribución dinámica en caché sobre la lógica de seguimiento de gráficos de TensorFlow. Para ser más específico sobre la terminología:

  • Un tf.Graph es la representación portátil sin procesar, tf.Graph del lenguaje, de un cálculo de TensorFlow.
  • Un ConcreteFunction envuelve un tf.Graph .
  • Una Function administra una caché de ConcreteFunction sy elige la correcta para sus entradas.
  • tf.function envuelve una función de Python, devolviendo un objeto Function .
  • El rastreo crea un tf.Graph y lo envuelve en una función ConcreteFunction , también conocida como rastreo.

Reglas de rastreo

Una Function determina si reutilizar una función ConcreteFunction rastreada calculando una clave de caché a partir de los args y kwargs de una entrada. Una clave de caché es una clave que identifica una función ConcreteFunction función de los argumentos y kwargs de entrada de la llamada a la Function , de acuerdo con las siguientes reglas (que pueden cambiar):

  • La clave generada para un tf.Tensor es su forma y dtype.
  • La clave generada para una tf.Variable es una identificación de variable única.
  • La clave generada para una primitiva de Python (como int , float , str ) es su valor.
  • La clave generada para dict s anidados, list s, tuple s, namedtuple s y attr s es la tupla aplanada de leaf-keys (vernest.flatten ). (Como resultado de este aplanamiento, llamar a una función concreta con una estructura de anidamiento diferente a la utilizada durante el rastreo dará como resultado un TypeError).
  • Para todos los demás tipos de Python, la clave es única para el objeto. De esta manera, se rastrea una función o método de forma independiente para cada instancia con la que se llama.

Controlar el retroceso

Retrazar, que es cuando su Function crea más de un seguimiento, ayuda a garantizar que TensorFlow genere gráficos correctos para cada conjunto de entradas. Sin embargo, el rastreo es una operación costosa. Si su Function traza un nuevo gráfico para cada llamada, encontrará que su código se ejecuta más lentamente que si no usara tf.function .

Para controlar el comportamiento de seguimiento, puede utilizar las siguientes técnicas:

  • Especifique input_signature en tf.function para limitar el seguimiento.
@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 una dimensión [None] en tf.TensorSpec para permitir flexibilidad en la reutilización de trazas.

    Dado que TensorFlow hace coincidir los tensores en función de su forma, el uso de una dimensión None como comodín permitirá que Function s reutilice trazas para entradas de tamaño variable. La entrada de tamaño variable puede ocurrir si tiene secuencias de diferente longitud o imágenes de diferentes tamaños para cada lote (consulte los tutoriales de Transformer y Deep Dream, por ejemplo).

@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)
  • Transmita argumentos de Python a tensores para reducir el retroceso.

    A menudo, los argumentos de Python se utilizan para controlar hiperparámetros y construcciones de gráficos, por ejemplo, num_layers=10 o training=True o nonlinearity='relu' . Entonces, si el argumento de Python cambia, tiene sentido que tenga que volver sobre el gráfico.

    Sin embargo, es posible que no se esté utilizando un argumento de Python para controlar la construcción del gráfico. En estos casos, un cambio en el valor de Python puede desencadenar un retroceso innecesario. Tomemos, por ejemplo, este ciclo de entrenamiento, que AutoGraph desenrollará dinámicamente. A pesar de las múltiples trazas, el gráfico generado es en realidad idéntico, por lo que no es necesario volver a rastrearlo.

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

Si necesita forzar el retroceso, cree una nueva Function . Se garantiza que los objetos de Function separados no compartirán trazas.

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

tf.function(f)()
tf.function(f)()
Tracing!
Executing
Tracing!
Executing

Obtener funciones concretas

Cada vez que se rastrea una función, se crea una nueva función concreta. Puede obtener directamente una función concreta, utilizando 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)

La impresión de una función ConcreteFunction muestra un resumen de sus argumentos de entrada (con tipos) y su tipo de salida.

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

También puede recuperar directamente la firma de una función 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)

El uso de una traza concreta con tipos incompatibles arrojará un error

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]

Puede notar que los argumentos de Python reciben un tratamiento especial en la firma de entrada de una función concreta. Antes de TensorFlow 2.3, los argumentos de Python simplemente se eliminaban de la firma de la función concreta. A partir de TensorFlow 2.3, los argumentos de Python permanecen en la firma, pero están limitados a tomar el valor establecido durante el seguimiento.

@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

Obtención de gráficos

Cada función concreta es una envoltura invocable alrededor de un tf.Graph . Aunque recuperar el objeto tf.Graph real no es algo que normalmente necesitará hacer, puede obtenerlo fácilmente desde cualquier función 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

Depuración

En general, depurar el código es más fácil en modo ansioso que dentro de tf.function . Debe asegurarse de que su código se ejecute sin errores en modo ansioso antes de decorar con tf.function . Para ayudar en el proceso de depuración, puede llamar a tf.config.run_functions_eagerly(True) para deshabilitar y volver a habilitar tf.function .

Al rastrear problemas que solo aparecen dentro de tf.function , aquí hay algunos consejos:

  • Las llamadas de print Python antiguas y sencillas solo se ejecutan durante el seguimiento, lo que le ayuda a rastrear cuándo se (re) rastrea su función.
  • tf.print llamadas a tf.print se ejecutarán cada vez y pueden ayudarlo a rastrear valores intermedios durante la ejecución.
  • tf.debugging.enable_check_numerics es una manera fácil de rastrear dónde se crean NaN e Inf.
  • pdb puede ayudarlo a comprender lo que sucede durante el rastreo. (Advertencia: PDB lo colocará en el código fuente transformado por AutoGraph).

Transformaciones de AutoGraph

AutoGraph es una biblioteca que está tf.function forma predeterminada en tf.function y transforma un subconjunto del código ansioso de Python en operaciones de TensorFlow compatibles con gráficos. Esto incluye el flujo de control como if , for , while .

Las operaciones de TensorFlow como tf.cond y tf.while_loop continúan funcionando, pero el flujo de control suele ser más fácil de escribir y comprender cuando se escribe en 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)>

Si tiene curiosidad, puede inspeccionar el código que genera el autógrafo.

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)

Condicionales

AutoGraph convertirá algunas declaraciones if <condition> en las llamadas tf.cond equivalentes. Esta sustitución se realiza si <condition> es un tensor. De lo contrario, la instrucción if se ejecuta como un condicional de Python.

Un condicional de Python se ejecuta durante el seguimiento, por lo que se agregará al gráfico exactamente una rama del condicional. Sin AutoGraph, este gráfico trazado no podría tomar la rama alternativa si hay un flujo de control dependiente de los datos.

tf.cond rastrea y agrega ambas ramas del condicional al gráfico, seleccionando dinámicamente una rama en el momento de la ejecución. El rastreo puede tener efectos secundarios no deseados; consulte Efectos de trazado de AutoGraph para obtener más información.

@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 la documentación de referencia para conocer las restricciones adicionales sobre las declaraciones if convertidas en AutoGraph.

Bucles

AutoGraph convertirá algunas declaraciones for y while en las operaciones de bucle de TensorFlow equivalentes, como tf.while_loop . tf.while_loop . Si no convertido, el for o while bucle se ejecuta como un bucle Python.

Esta sustitución se realiza en las siguientes situaciones:

Un bucle de Python se ejecuta durante el seguimiento, agregando operaciones adicionales al tf.Graph para cada iteración del bucle.

Un bucle de TensorFlow rastrea el cuerpo del bucle y selecciona dinámicamente cuántas iteraciones se ejecutarán en el momento de la ejecución. El cuerpo del bucle solo aparece una vez en el tf.Graph generado.

Consulte la documentación de referencia para conocer las restricciones adicionales sobre las declaraciones for y while convertidas en AutoGraph.

Bucle sobre datos de Python

Un error común es recorrer los datos de Python / Numpy dentro de una función tf.function . Este bucle se ejecutará durante el proceso de rastreo, agregando una copia de su modelo al tf.Graph para cada iteración del bucle.

Si desea envolver todo el ciclo de entrenamiento en tf.function , la forma más segura de hacerlo es envolver sus datos como untf.data.Dataset para que AutoGraph desenrolle dinámicamente el ciclo de entrenamiento.

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

Al envolver datos de Python / Numpy en un conjunto de datos, tenga en cuenta tf.data.Dataset.from_generator versus tf.data.Dataset.from_tensors . El primero mantendrá los datos en Python y los tf.py_function través de tf.py_function que puede tener implicaciones de rendimiento, mientras que el segundo tf.constant() una copia de los datos como un nodo tf.constant() grande en el gráfico, lo que puede tener implicaciones de memoria.

Leer datos de archivos a través de TFRecordDataset / CsvDataset / etc. es la forma más efectiva de consumir datos, ya que entonces TensorFlow mismo puede administrar la carga asincrónica y la captura previa de datos, sin tener que involucrar a Python. Para obtener más información, consulte la guía tf.data .

Acumulando valores en un bucle

Un patrón común es acumular valores intermedios de un bucle. Normalmente, esto se logra agregando a una lista de Python o agregando entradas a un diccionario de Python. Sin embargo, como se trata de efectos secundarios de Python, no funcionarán como se esperaba en un ciclo desenrollado dinámicamente. Utilice tf.TensorArray para acumular resultados de un bucle desenrollado dinámicamente.

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

Limitaciones

La Function TensorFlow tiene algunas limitaciones de diseño que debe tener en cuenta al convertir una función de Python en una Function .

Ejecutando efectos secundarios de Python

Los efectos secundarios, como imprimir, agregar a listas y cambiar globales, pueden comportarse inesperadamente dentro de una Function , a veces ejecutándose dos veces o no todas. Solo suceden la primera vez que llama a una Function con un conjunto de entradas. Posteriormente, el tf.Graph trazado se vuelve a ejecutar, sin ejecutar el código Python.

La regla general es evitar depender de los efectos secundarios de Python en su lógica y usarlos solo para depurar sus rastros. De lo contrario, las API de TensorFlow como tf.data , tf.print , tf.summary , tf.Variable.assign y tf.TensorArray son la mejor manera de garantizar que el tiempo de ejecución de TensorFlow ejecute tu código con cada llamada.

@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

Si desea ejecutar código Python durante cada invocación de una Function , tf.py_function es una trampilla de salida. El inconveniente de tf.py_function es que no es portátil ni tiene un rendimiento particular, no se puede guardar con SavedModel y no funciona bien en configuraciones distribuidas (multi-GPU, TPU). Además, dado que tf.py_function debe conectarse al gráfico, envía todas las entradas / salidas a los tensores.

Cambiar las variables globales y libres de Python

Cambiar las variables globales y libres de Python cuenta como un efecto secundario de Python, por lo que solo ocurre durante el seguimiento.

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

Debe evitar la mutación de contenedores como listas, dictados, otros objetos que viven fuera de la Function . En su lugar, use argumentos y objetos TF. Por ejemplo, la sección "Acumulación de valores en un bucle" tiene un ejemplo de cómo se pueden implementar operaciones de tipo lista.

En algunos casos, puede capturar y manipular el estado si es una tf.Variable . Así es como se actualizan los pesos de los modelos de Keras con llamadas repetidas a la misma función ConcreteFunction .

Usando iteradores y generadores de Python

Muchas características de Python, como generadores e iteradores, dependen del tiempo de ejecución de Python para realizar un seguimiento del estado. En general, si bien estas construcciones funcionan como se esperaba en el modo ansioso, son ejemplos de efectos secundarios de Python y, por lo tanto, solo ocurren durante el seguimiento.

@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

Al igual que TensorFlow tiene un tf.TensorArray especializado para construcciones de listas, tiene un tf.data.Iterator especializado para construcciones de iteración. Consulte la sección sobre Transformaciones de AutoGraph para obtener una descripción general. Además, la API tf.data puede ayudar a implementar patrones de generador:

@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

Eliminando tf.Variables entre llamadas a Function

Otro error que puede encontrar es una variable de recolección de basura. ConcreteFunction s solo retiene WeakRefs para las variables sobre las que se cierran, por lo que debe conservar una referencia a cualquier variable.

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 conocidos

Si su Function no se evalúa correctamente, el error puede deberse a estos problemas conocidos que se planea solucionar en el futuro.

Dependiendo de las variables globales y libres de Python

Function crea una nueva función ConcreteFunction cuando se llama con un nuevo valor de un argumento de Python. Sin embargo, no hace eso para el cierre de Python, globales o no locales de esa Function . Si su valor cambia entre llamadas a la Function , la Function seguirá usando los valores que tenía cuando se trazó. Esto es diferente de cómo funcionan las funciones normales de Python.

Por esa razón, recomendamos un estilo de programación funcional que use argumentos en lugar de cerrar sobre nombres 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)

Puede cerrar los nombres externos, siempre que no actualice sus valores.

Dependiendo de los objetos de Python

La recomendación de pasar objetos de Python como argumentos a tf.function tiene una serie de problemas conocidos que se espera que se solucionen en el futuro. En general, puede confiar en un seguimiento consistente si usa una estructura primitiva de Python o compatible con tf.nest como argumento o pasa una instancia diferente de un objeto a una Function . Sin embargo, Function no creará una nueva traza cuando pase el mismo objeto y solo cambie sus 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)

El uso de la misma Function para evaluar la instancia actualizada del modelo tendrá errores, ya que el modelo actualizado tiene la misma clave de caché que el modelo original.

Por esa razón, le recomendamos que escriba su Function para evitar depender de atributos de objetos mutables o crear nuevos objetos.

Si eso no es posible, una solución es crear nuevas Function cada vez que modifique su objeto para forzar el retroceso:

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 el retroceso puede ser costoso , puede usar tf.Variable s como atributos de objeto, que se pueden mutar (pero no cambiar, ¡cuidado!) Para obtener un efecto similar sin necesidad de retroceder.

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)

Creando tf.Variables

Function solo admite la creación de variables una vez, cuando se llama por primera vez, y luego reutilizarlas. No puede crear tf.Variables en nuevos seguimientos. La creación de nuevas variables en llamadas posteriores no está permitida actualmente, pero lo estará en el futuro.

Ejemplo:

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

Puede crear variables dentro de una Function siempre que esas variables solo se creen la primera vez que se ejecuta la función.

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)

Usar con múltiples optimizadores de Keras

Puede encontrar ValueError: tf.function-decorated function tried to create variables on non-first call. cuando se utiliza más de un optimizador de Keras con una función tf.function . Este error se produce porque los optimizadores crean internamente tf.Variables cuando aplican gradientes por primera 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.

Si necesita cambiar el optimizador durante el entrenamiento, una solución es crear una nueva Function para cada optimizador, llamando a ConcreteFunction directamente.

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.

Usar con múltiples modelos de Keras

También puede encontrar ValueError: tf.function-decorated function tried to create variables on non-first call. al pasar diferentes instancias de modelo a la misma Function .

Este error se produce porque los modelos de Keras (que no tienen definida su forma de entrada ) y las capas de Keras crean tf.Variables s cuando se llaman por primera vez. Puede estar intentando inicializar esas variables dentro de una Function , que ya ha sido llamada. Para evitar este error, intente llamar a model.build(input_shape) para inicializar todos los pesos antes de entrenar el modelo.

Otras lecturas

Para obtener información sobre cómo exportar y cargar una Function , consulte la guía de modelo guardado . Para obtener más información sobre las optimizaciones de gráficos que se realizan después del seguimiento, consulte la guía de Grappler . Para saber cómo optimizar su canalización de datos y perfilar su modelo, consulte la guía Profiler .