¿Debo usar @ tf.function para todas las funciones?

Un tutorial oficial sobre @tf.functiondice:

Para obtener el máximo rendimiento y hacer que su modelo sea desplegable en cualquier lugar, use tf.function para hacer gráficos de sus programas. Gracias a AutoGraph, una sorprendente cantidad de código de Python simplemente funciona con tf.function, pero todavía hay dudas de las que debemos tener cuidado.

Las principales conclusiones y recomendaciones son:

• No confíe en los efectos secundarios de Python como la mutación de objetos o los anexos de listas.
• La función tf funciona mejor con operaciones TensorFlow, en lugar de operaciones NumPy o primitivas Python.
• En caso de duda, utilice el idioma for x in y.

Solo menciona cómo implementar @tf.functionfunciones anotadas, pero no cuándo usarlo.

¿Hay una heurística sobre cómo decidir si al menos debería intentar anotar una función tf.function? Parece que no hay razones para no hacerlo, a menos que sea demasiado flojo para eliminar los efectos secundarios o cambiar algunas cosas como range()-> tf.range(). Pero si estoy dispuesto a hacer esto ...

¿Hay alguna razón para no usar @tf.functionpara todas las funciones?

TLDR: depende de su función y de si está en producción o desarrollo. No lo use tf.functionsi desea poder depurar su función fácilmente, o si se encuentra bajo las limitaciones de AutoGraph o la compatibilidad del código tf.v1. Recomiendo ver Inside TensorFlow habla sobre AutoGraph y Funciones, no sobre Sesiones .

A continuación, desglosaré los motivos, todos tomados de la información disponible en línea por Google.

En general, el tf.functiondecorador hace que una función se compile como invocable que ejecuta un gráfico TensorFlow. Esto implica:

• Conversión del código a través de AutoGraph si es necesario (incluidas las funciones llamadas desde una función anotada)
• Rastreo y ejecución del código gráfico generado

Hay información detallada disponible sobre las ideas de diseño detrás de esto.

Beneficios de decorar una función con tf.function

Beneficios generales

• Ejecución más rápida , especialmente si la función consta de muchas operaciones pequeñas (Fuente)

Para funciones con código Python / Uso de AutoGraph vía tf.function decoración

Si desea usar AutoGraph, tf.functionse recomienda usarlo en lugar de llamar directamente a AutoGraph. Las razones para esto incluyen: dependencias de control automático, se requiere para algunas API, más almacenamiento en caché y ayudantes de excepción (Fuente) .

Inconvenientes de decorar una función con tf.function

Inconvenientes generales

• Si la función solo consta de pocas operaciones costosas, no habrá mucha aceleración (Fuente)

Para funciones con código Python / Uso de AutoGraph mediantetf.function decoración

• Captura sin excepción (debe hacerse en modo ansioso; fuera de la función decorada) (Fuente)
• La depuración es mucho más difícil.
• Limitaciones debido a efectos secundarios ocultos y flujo de control TF

La información detallada sobre las limitaciones de AutoGraph está disponible.

Para funciones con código tf.v1

• No está permitido crear variables más de una vez tf.function, pero está sujeto a cambios a medida que el código tf.v1 se elimina (Fuente)

Para funciones con código tf.v2

• No hay inconvenientes específicos

Ejemplos de limitaciones.

Crear variables más de una vez

No está permitido crear variables más de una vez, como ven el siguiente ejemplo:

@tf.function
def f(x):
    v = tf.Variable(1)
    return tf.add(x, v)

f(tf.constant(2))

# => ValueError: tf.function-decorated function tried to create variables on non-first call.

En el siguiente código, esto se mitiga asegurándose de que self.vsolo se crea una vez:

class C(object):
    def __init__(self):
        self.v = None
    @tf.function
    def f(self, x):
        if self.v is None:
            self.v = tf.Variable(1)
        return tf.add(x, self.v)

c = C()
print(c.f(tf.constant(2)))

# => tf.Tensor(3, shape=(), dtype=int32)

Efectos secundarios ocultos no capturados por AutoGraph

Los cambios como self.aen este ejemplo no se pueden ocultar, lo que conduce a un error ya que el análisis de funciones cruzadas aún no se realiza (Fuente) :

class C(object):
    def change_state(self):
        self.a += 1

    @tf.function
    def f(self):
        self.a = tf.constant(0)
        if tf.constant(True):
            self.change_state() # Mutation of self.a is hidden
        tf.print(self.a)

x = C()
x.f()

# => InaccessibleTensorError: The tensor 'Tensor("add:0", shape=(), dtype=int32)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=cond_true_5, id=5477800528); accessed from: FuncGraph(name=f, id=5476093776).

Los cambios a simple vista no son un problema:

class C(object):
    @tf.function
    def f(self):
        self.a = tf.constant(0)
        if tf.constant(True):
            self.a += 1 # Mutation of self.a is in plain sight
        tf.print(self.a)

x = C()
x.f()

# => 1

Ejemplo de limitación debido al flujo de control TF

Esta declaración if conduce a un error porque el valor de else debe definirse para el flujo de control TF:

@tf.function
def f(a, b):
    if tf.greater(a, b):
        return tf.constant(1)

# If a <= b would return None
x = f(tf.constant(3), tf.constant(2))   

# => ValueError: A value must also be returned from the else branch. If a value is returned from one branch of a conditional a value must be returned from all branches.

tf.function es útil en la creación y el uso de gráficos computacionales, deben usarse en el entrenamiento y en la implementación, sin embargo, no es necesario para la mayoría de sus funciones.

Digamos que estamos construyendo una capa especial que estará separada de un modelo más grande. No queremos tener el decorador tf.function encima de la función que construye esa capa porque es simplemente una definición de cómo se verá la capa.

Por otro lado, digamos que vamos a hacer una predicción o continuar nuestro entrenamiento usando alguna función. Nos gustaría tener el decorador tf.function porque en realidad estamos usando el gráfico computacional para obtener algún valor.

Un gran ejemplo sería construir un modelo codificador-decodificador. NO coloque el decorador alrededor de la función para crear el codificador o decodificador o cualquier capa, eso es solo una definición de lo que hará. SÍ coloque el decorador alrededor del método "entrenar" o "predecir" porque esos realmente van a usar el gráfico computacional para el cálculo.