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.

Mejor rendimiento con tf.function

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar bloc de notas

En TensorFlow 2, la ejecución eager 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 ir en detrimento del rendimiento y la capacidad de implementación.

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

Esta guía le ayudará a conceptualizar cómo funciona tf.function por dentro, para que pueda usarlo con eficacia.

Las principales conclusiones y recomendaciones son:

• Depure en modo eager, luego decore con @tf.function.

• No dependa de los efectos secundarios de Python, como la mutación de objetos o la anexión de listas.

• tf.function funciona mejor con las ops de TensorFlow; las llamadas de NumPy y Python se convierten en constantes.

Preparación

import tensorflow as tf

Defina una función ayudante 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))

Conceptos básicos

Uso

Una Function que usted defina (por ejemplo, aplicando el decorador @tf.function) es igual que una operación central de TensorFlow: Puede ejecutarse de forma eager; puede calcular gradientes; etc.

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

Puede usar Functions dentro de otras Functions.

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

Functions puede ser más rápido que el código eager, especialmente para grafos con muchas ops pequeñas. Pero para grafos con unas pocas ops costosas (como las convoluciones), es posible que no se vea mucho aumento de la velocidad.

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

Trazado

Esta sección expone cómo Función funciona por dentro, incluyendo detalles de implementación que pueden cambiar en el futuro. De todos modos, en cuanto entienda por qué y cuándo se produce el trazado, ¡le resultará mucho más fácil usar tf.function con eficacia!

¿Qué es el "trazado"?

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

Function cubre este vacío al separar su código en dos etapas:

1. En la primera etapa, llamada "trazado", Function crea un nuevo tf.Graph. El código Python se ejecuta normalmente, pero todas las operaciones TensorFlow (como sumar dos Tensores) están aplazadas: son capturadas por el tf.Graph y no se ejecutan.

2. En la segunda etapa, se ejecuta un tf.Graph que contiene todo lo que se aplazó en la primera etapa. Esta etapa es mucho más rápida que la etapa de trazado.

Dependiendo de sus entradas, Function no siempre ejecutará la primera etapa cuando sea llamada. Vea más abajo las "Reglas de trazado" si quiere entender mejor cómo toma esa determinación. TensorFlow tiene un alto rendimiento porque se salta la primera etapa y sólo ejecuta la segunda.

Cuando Function sí decide trazar, la etapa de trazado es seguida inmediatamente por la segunda etapa, por lo que llamar a Function tanto crea como ejecuta el tf.Graph. Más adelante verá cómo puede ejecutar sólo la etapa de trazado con get_concrete_function.

Cuando se pasan argumentos de distinto tipo 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()

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

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

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

print(double.pretty_printed_concrete_signatures())

Hasta ahora, ha visto que tf.function crea una capa de envío dinámica y en caché sobre la lógica de trazado del grafo de TensorFlow. Para ser más específico sobre la terminología:

• Un tf.Graph es la representación en bruto, agnóstica al lenguaje y portátil de un cálculo TensorFlow.

• Una ConcreteFunction encapsula un tf.Graph.

• Una Function administra una caché de ConcreteFunctions y elige la adecuada para sus entradas.

• tf.function encapsula una función Python, devolviendo un objeto Function.

• El trazado crea un tf.Graph y lo encapsula en una ConcreteFunction, también conocida como un trazo.

Reglas de trazado

Cuando se llama, una Function hace coincidir los argumentos de llamada con ConcreteFunctions existentes usando tf.types.experimental.TraceType de cada argumento. Si se encuentra una ConcreteFunction que coincida, se le envía la llamada. Si no se encuentra ninguna coincidencia, se traza una nueva ConcreteFunction.

Si se encuentran varias coincidencias, se selecciona la firma más específica. La coincidencia se realiza mediante subtipado, de forma muy parecida a las llamadas a funciones normales en C++ o Java, por ejemplo. Por ejemplo, TensorShape([1, 2]) es un subtipo de TensorShape([None, None]) y, por tanto, una llamada a tf.function con TensorShape([1, 2]) puede ser enviada a la ConcreteFunction producida con TensorShape([None, None]) pero si también existe una ConcreteFunction con TensorShape([1, None]) entonces se le dará prioridad ya que es más específica.

El TraceType se determina a partir de los argumentos de entrada como sigue:

• Para Tensor, el tipo está parametrizado por los dtype y shape del Tensor; las formas clasificadas son un subtipo de las formas no clasificadas; las dimensiones fijas son un subtipo de las dimensiones desconocidas

• Para Variable, el tipo es similar a Tensor, pero también incluye un ID de recurso único de la variable, necesario para integrar correctamente las dependencias de control.

• Para los valores primitivos de Python, el tipo corresponde al propio valor. Por ejemplo, el TraceType del valor 3 es LiteralTraceType<3>, no int.

• Para los contenedores ordenados de Python como list y tuple, etc., el tipo está parametrizado por los tipos de sus elementos; por ejemplo, el tipo de [1, 2] es ListTraceType<LiteralTraceType<1>, LiteralTraceType<2>> y el tipo para [2, 1] es ListTraceType<LiteralTraceType<2>, LiteralTraceType<1>> que es diferente.

• Para mapeos en Python como dict, el tipo es también un mapeo de las mismas claves pero a los tipos de valores en lugar de a los valores reales. Por ejemplo, el tipo de {1: 2, 3: 4}, es MappingTraceType<<KeyValue<1, LiteralTraceType<2>>>, <KeyValue<3, LiteralTraceType<4>>>>. Sin embargo, a diferencia de los contenedores ordenados, {1: 2, 3: 4} y {3: 4, 1: 2} tienen tipos equivalentes.

• Para los objetos Python que implementan el método __tf_tracing_type__, el tipo es cualquier cosa que devuelva ese método

• Para cualquier otro objeto Python, el tipo es un TraceType genérico, su precedencia de coincidencia es:

  • Primero comprueba si el objeto es el mismo que se usó en el trazo anterior (usando python id() o is). Tenga en cuenta que esto seguirá coincidiendo si el objeto ha cambiado, por lo que si usa objetos python como argumentos tf.function es mejor usar los inmutables.

  • A continuación comprueba si el objeto es igual al objeto usado en el trazo anterior (usando python ==).

  Tenga en cuenta que este procedimiento sólo conserva una referencia débil (weakref) al objeto y, por lo tanto, sólo funciona mientras el objeto esté en el ámbito de aplicación/no se haya eliminado).

Nota: TraceType se basa en los parámetros de entrada Function por lo que los cambios en las variables globales y libres por sí solos no crearán un nuevo trazo. Consulte esta sección para conocer las prácticas recomendadas al tratar con variables globales y libres de Python.

Control del retrazado

El retrazado, que es cuando su Function crea más de un trazo, ayuda a asegurar que TensorFlow genera grafos correctos para cada conjunto de entradas. Sin embargo, ¡el trazado es una operación costosa! Si su Function retraza un nuevo grafo para cada llamada, se dará cuenta de que su código se ejecuta más lentamente que si no usara tf.function.

Para controlar el comportamiento del trazado, puede usar las siguientes técnicas:

Pasar una input_signature fija a 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]))

Usar dimensiones desconocidas para mayor flexibilidad

Ya que TensorFlow hace coincidir los tensores basándose en su forma, usar una dimensión None como comodín permitirá a las Functions reutilizar trazos 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 (Vea los tutoriales 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])))

Pasar tensores en lugar de literales python

Los argumentos de Python se usan a menudo para controlar los hiperparámetros y la construcción de grafos; por ejemplo, num_layers=10 o training=True o nonlinearity='relu'. Por lo tanto, si el argumento Python cambia, es lógico que tenga que retrazar el grafo.

Sin embargo, podría darse el caso de que no se esté usando un argumento de Python para controlar la construcción del grafo. En estos casos, un cambio en el valor de Python puede desencadenar un retrazado innecesario. Tomemos, por ejemplo, este bucle de entrenamiento, que AutoGraph desenrollará dinámicamente. A pesar de las múltiples trazas, el grafo generado es en realidad idéntico, por lo que el retrazado es innecesario.

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

Si necesita forzar el retrazado, cree una nueva Function. Se garantiza que los objetos Function separados no comparten trazos.

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

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

Usar el protocolo de trazado

Siempre que sea posible, debería preferir en su lugar convertir el tipo Python en un tf.experimental.ExtensionType. Además, el TraceType de un ExtensionType es el tf.TypeSpec asociado a él. Por lo tanto, si es necesario, puede simplemente anular el tf.TypeSpec predeterminado para tomar el control del Tracing Protocol de un ExtensionType. Para más detalles, consulte la sección Personalizar el TypeSpec de ExtensionType en la guía Tipos de extensión.

De lo contrario, para tener un control directo sobre cuándo Function debe retrazar con respecto a un tipo particular de Python, puede implementar usted mismo el Tracing Protocol para ello.

@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

Obtener funciones concretas

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

Al imprimir una ConcreteFunction se muestra un resumen de sus argumentos de entrada (con tipos) y su tipo de salida.

print(double_strings)

También puede recuperar directamente la firma de una función concreta.

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

Usar un trazo concreto con tipos incompatibles arrojará un error

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

Puede que haya notado que los argumentos Python reciben un tratamiento especial en la firma de entrada de una función concreta. Antes de TensorFlow 2.3, los argumentos Python eran simplemente eliminados de la firma de la función concreta. A partir de TensorFlow 2.3, los argumentos Python permanecen en la firma, pero están restringidos a tomar el valor fijado durante el trazado.

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

Obtener grafos

Cada función concreta es un contenedor invocable que envuelve un tf.Graph. Aunque recuperar el objeto tf.Graph real no es algo que necesite hacer normalmente, puede hacerlo 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}')

Depuración

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

A continuación le ofrecemos algunos consejos para el seguimiento de los problemas que sólo aparecen dentro de tf.function:

• Las llamadas simples print de Python sólo se ejecutan durante el trazado, lo que le ayuda a realizar un seguimiento de cuándo se (re)traza su función.

• Las llamadas tf.print se ejecutarán cada vez, y pueden ayudarle a realizar un seguimiento de los valores intermedios durante la ejecución.

• tf.debugging.enable_check_numerics es una forma sencilla de descubrir dónde se crean los NaN y los Inf.

• pdb (el depurador de Python) puede ayudarle a entender lo que ocurre durante el trazado. (Advertencia: pdb le llevará al código fuente transformado por AutoGraph).

Transformaciones de AutoGraph

AutoGraph es una librería que está activada de forma predeterminada en tf.function, y transforma un subgrupo de código eager de Python en ops de TensorFlow compatibles con grafos. Esto incluye flujos de control como if, for, while.

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

Si le interesa, puede inspeccionar el código que genera el autograph.

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

Condicionales

AutoGraph convertirá algunas sentencias if <condition> en las llamadas tf.cond equivalentes. Esta sustitución se realiza si <condition> es un Tensor. En caso contrario, la sentencia if se ejecuta como un condicional de Python.

Un condicional Python se ejecuta durante el trazado, por lo que se añadirá exactamente una derivación del condicional al grafo. Sin AutoGraph, este grafo trazado sería incapaz de tomar la bifurcación alternativa si existe un flujo de control dependiente de los datos.

tf.cond traza y añade ambas bifurcaciones del condicional al grafo, seleccionando dinámicamente una derivación en el momento de la ejecución. El trazado puede tener efectos secundarios no deseados; revise efectos del trazado de AutoGraph si desea 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))

Si desea saber más sobre las restricciones de las sentencias if convertidas con AutoGraph, consulte la documentación de referencia.

Bucles

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

Esta sustitución se realiza en las siguientes situaciones:

• for x in y: si y es un Tensor, convierta a tf.while_loop. En el caso especial de que y sea un tf.data.Dataset, se genera una combinación de ops tf.data.Dataset.

• while <condición>: si <condition> es un tensor, convierta a tf.while_loop.

Un bucle Python se ejecuta durante el trazado, añadiendo ops adicionales al tf.Graph por cada iteración del bucle.

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

Si desea saber más sobre las restricciones de las sentencias for y while convertidas con AutoGraph, consulte la documentación de referencia.

Bucle sobre datos Python

Un error común es hacer un bucle sobre los datos de Python/NumPy dentro de una tf.function. Este bucle se ejecutará durante el proceso de trazado, añadiendo una copia de su modelo al tf.Graph por cada iteración del bucle.

Si desea encapsular todo el bucle de entrenamiento en tf.function, la forma más segura de hacerlo es encapsular sus datos como un tf.data.Dataset para que AutoGraph desenrede dinámicamente el bucle 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)))

Cuando encapsule datos Python/NumPy en un conjunto de datos, preste atención a tf.data.Dataset.from_generator versus tf.data.Dataset.from_tensor_slices. El primero conservará los datos en Python y los recuperará mediante tf.py_function, lo que puede tener implicaciones en el rendimiento, mientras que el segundo agrupará una copia de los datos como un gran nodo tf.constant() en el grafo, lo que puede tener implicaciones en la memoria.

La lectura de datos de archivos a través de TFRecordDataset, CsvDataset, etc. es la forma más eficaz de consumir datos, ya que entonces el propio TensorFlow puede administrar la carga asíncrona y la preextracción de datos, sin tener que involucrar a Python. Si desea obtener más información, consulte el tf.data: Construir canalizaciones de entrada de TensorFlow.

Valores acumulados en un bucle

Un patrón común es acumular valores intermedios de un bucle. Normalmente, esto se consigue añadiendo a una lista Python o añadiendo entradas a un diccionario Python. Sin embargo, como estos son efectos secundarios de Python, no funcionarán como se espera en un bucle desenredado dinámicamente. Use tf.TensorArray para acumular resultados de un bucle desenredado 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]))

Limitaciones

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

Efectos secundarios de la ejecución de Python

Los efectos secundarios, como imprimir, anexar a listas y mutar globales, pueden comportarse de forma inesperada dentro de una Function, a veces ejecutándose dos veces o ninguna. Sólo ocurren la primera vez que se llama a una Function con un conjunto de entradas. Después, se vuelve a ejecutar el tf. Graph trazado, sin ejecutar el código Python.

La regla general es evitar depender de los efectos secundarios de Python en su lógica y sólo usarlos para depurar sus trazados. De lo contrario, la mejor forma de asegurarse de que su código será ejecutado por el runtime de TensorFlow con cada llamada son las APIs de TensorFlow como tf.data, tf.print, tf.summary, tf.Variable.assign, y tf.TensorArray.

@tf.function
def f(x):
  print("Traced with", x)
  tf.print("Executed with", x)

f(1)
f(1)
f(2)

Si desea ejecutar código Python durante cada invocación de una Function, tf.py_function es una solución de salida. El inconveniente de tf.py_function es que no es portable ni especialmente eficiente, no puede guardarse con SavedModel y no funciona bien en configuraciones distribuidas (multi-GPU, TPU). Además, dado que tf.py_function tiene que ser integrada en el grafo, convierte todas las entradas/salidas en tensores.

Cambio de variables globales y libres de Python

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

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

A veces, los comportamientos inesperados son muy difíciles de notar. En el ejemplo siguiente, se pretende que el counter garantice el incremento de una variable. Sin embargo, debido a que es un entero python y no un objeto TensorFlow, su valor es capturado durante el primer trazado. Cuando se usa tf.function, assign_add se graba incondicionalmente en el grafo subyacente. Por lo tanto, v aumentará en 1, cada vez que se llame a la tf.function. Este problema es común entre los usuarios que intentan migrar su código Tensorflow en modo Graph a Tensorflow 2 usando decoradores tf.function, cuando los efectos secundarios de Python (counter en el ejemplo) se usan para determinar qué ops ejecutar (assign_add en el ejemplo). Normalmente, los usuarios se dan cuenta de esto sólo después de ver resultados numéricos sospechosos, o un rendimiento significativamente inferior al esperado (por ejemplo, si la operación vigilada es muy costosa).

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

Una solución para conseguir el comportamiento esperado es usar tf.init_scope para elevar las operaciones fuera del grafo de funciones. Esto garantiza que el incremento de la variable sólo se realice una vez durante el tiempo de trazado. Debe tenerse en cuenta que init_scope tiene otros efectos secundarios, incluyendo el flujo de control despejado y la cinta de gradiente. A veces el uso de init_scope puede llegar a ser demasiado complejo para administrarlo 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

En resumen, como regla general, debe evitar mutar objetos python, como enteros, o contenedores, como listas, que se encuentren fuera de la Function. En su lugar, use argumentos y objetos TF. Por ejemplo, la sección "Valores acumuladors en un bucle" tiene un ejemplo de cómo implementar operaciones tipo lista.

Puede, en algunos casos, capturar y manipular el estado si es una tf.Variable. Así es como se actualizan las ponderaciones de los modelos Keras con llamadas repetidas a la misma ConcreteFunction.

Usar iteradores y generadores de Python

Muchas funciones de Python, como los generadores y los iteradores, dependen de que el runtime de Python lleve el control de su estado. En general, aunque estos constructos funcionan como se espera en modo eager, son ejemplos de efectos secundarios de Python y, por tanto, sólo se producen durante el trazado.

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

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

@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 las salidas de una función tf.deben ser valores retornados

Con la excepción de tf.Variables, una función tf.debe retornar todas sus salidas. Intentar acceder directamente a cualquier tensor desde una función sin pasar por los valores de retorno provoca "fugas".

Por ejemplo, la función siguiente "fuga" el tensor a a través de la x global de 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)

Esto es verdadero incluso si también se retorna el valor filtrado:

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

Normalmente, este tipo de fugas se producen cuando se usan sentencias o estructuras de datos de Python. Además de filtrar tensores inaccesibles, tales declaraciones también son probablemente erróneas porque cuentan como efectos secundarios de Python, y no se garantiza que se ejecuten en cada llamada de función.

Entre las formas habituales de fugas de tensores locales también se incluye mutar una colección externa de Python, o un 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

Las tf.functions recursivas no son compatibles

Functions recursivas no son compatibles y podrían provocar bucles infinitos. Por ejemplo,

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

Incluso si una Function recursiva parece funcionar, la función python será trazada varias veces y podría tener consecuencias para el rendimiento. Por ejemplo,

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

Si su Function no se evalúa correctamente, el error puede explicarse debido a estos problemas conocidos que está previsto solucionar en el futuro.

Depender de variables globales y libres de Python

Function crea una nueva ConcreteFunction cuando se llama con un nuevo valor de un argumento Python. Sin embargo, no lo hace para el cierre de Python, ni para sus globales o no locales de esa Function. Si su valor cambia entre llamadas a la Function, la Function seguirá usando los valores que tenían cuando fue trazada. Esto es diferente de como trabajan las funciones normales de Python.

Por este motivo, debe seguir un estilo de programación funcional que use argumentos en lugar de cierre 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))

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

Otra forma de actualizar un valor global, es convertirlo en una tf.Variable y usar en su lugar el 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

Es permitido pasar objetos Python personalizados como argumentos a tf.function, pero tiene ciertas limitaciones.

Para obtener la máxima cobertura de funciones, considere la posibilidad de transformar los objetos en Tipos de extensión antes de pasarlos a tf.function. También puede usar primitivas de Python y estructuras compatibles con tf.nest.

Sin embargo, como se explica en las reglas de trazado, cuando la clase personalizada de Python no ofrece un TraceType personalizado, tf.function se ve obligada a usar la igualdad basada en instancias, lo que significa que no creará un nuevo trazado cuando le pase el mismo objeto con 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 la misma Function para evaluar la instancia modificada del modelo tendrá errores, ya que sigue teniendo el mismo TraceType basado en instancias que el modelo original.

Por este motivo, se recomienda que escriba su Function para evitar depender de los atributos mutables de los objetos o que implemente el Protocolo de trazado para que los objetos informen a su Function sobre dichos atributos.

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

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 el retrazado puede resultar costoso, puede usar tf.Variables como atributos de objeto, que pueden mutarse (¡ojo! pero no cambiarse) para conseguir un efecto similar sin necesidad de retrazado.

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!

Crear tf.Variables

Function sólo admite tf.Variables únicos creados una vez en la primera llamada y reutilizados en las siguientes llamadas a la función. El siguiente fragmento de código crearía una nueva tf.Variable en cada llamada a la función, lo que resultaría en una excepción ValueError.

Ejemplo:

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

with assert_raises(ValueError):
  f(1.0)

Un patrón común utilizado para solucionar esta limitación es comenzar con un valor None de Python y, a continuación, crear condicionalmente la tf.Variable si el valor es 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())

Usar con múltiples optimizadores Keras

Puede encontrarse con el error ValueError: tf.function only supports singleton tf.Variables created on the first call. al usar más de un optimizador Keras con una 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)

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

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)

Usar múltiples modelos Keras

También puede encontrarse con el error ValueError: tf.function only supports singleton tf.Variables created on the first call. al pasar diferentes instancias del modelo a la misma Function.

Este error se produce porque los modelos Keras (que no tienen definida su forma de entrada) y las capas Keras crean tf.Variabless cuando son llamadas por primera vez. Puede que esté 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 todas las ponderaciones antes de entrenar el modelo.

Lecturas adicionales

Si desea más información sobre cómo exportar y cargar una Function, consulte la guía de SavedModel. Si desea más información sobre las optimizaciones de grafos que se realizan tras el trazado, consulte la guía de Grappler. Si desea saber cómo optimizar su canalización de datos y perfilar su modelo, consulte la guía de Profiler.