Descripción del decorador de funciones @tf. de TensorFlow

Introducción

Mejorar el rendimiento de un ciclo de entrenamiento puede ahorrar horas de tiempo de cómputo cuando se entrenan modelos de aprendizaje automático. Una de las formas de mejorar el rendimiento del código de TensorFlow es usar el tf.function() decorador: un cambio simple de una línea que puede hacer que sus funciones se ejecuten significativamente más rápido.

En esta breve guía, explicaremos cómo tf.function() mejora el rendimiento y echa un vistazo a algunas de las mejores prácticas.

Decoradores Python y tf.función()

En Python, un decorador es una función que modifica el comportamiento de otras funciones. Por ejemplo, suponga que llama a la siguiente función en una celda de cuaderno:

import tensorflow as tf

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

def some_costly_computation(x):
    aux = tf.eye(100, dtype=tf.dtypes.float32)
    result = tf.zeros(100, dtype = tf.dtypes.float32)
    for i in range(1,100):
        aux = tf.matmul(x,aux)/i
        result = result + aux
    return result

%timeit some_costly_computation(x)

16.2 ms ± 103 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Sin embargo, si pasamos la función costosa a una tf.function():

quicker_computation = tf.function(some_costly_computation)
%timeit quicker_computation(x)

Obtenemos quicker_computation() – una nueva función que se ejecuta mucho más rápido que la anterior:

4.99 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

¿Entonces tf.function() modifica some_costly_computation() y emite el quicker_computation() función. Los decoradores también modifican funciones, por lo que era natural hacer tf.function() un decorador también.

Usar la notación de decorador es lo mismo que llamar tf.function(function):

@tf.function
def quick_computation(x):
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

%timeit quick_computation(x)

5.09 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Cómo funciona tf.function() ¿Trabajo?

¿Cómo es que podemos hacer que ciertas funciones se ejecuten 2-3 veces más rápido?

El código de TensorFlow se puede ejecutar en dos modos: modo ansioso y modo gráfico. El modo Eager es la forma estándar e interactiva de ejecutar código: cada vez que llamas a una función, se ejecuta.

El modo gráfico, sin embargo, es un poco diferente. En modo gráfico, antes de ejecutar la función, TensorFlow crea un gráfico de cálculo, que es una estructura de datos que contiene las operaciones necesarias para ejecutar la función. El gráfico de cálculo permite que TensorFlow simplifique los cálculos y encuentre oportunidades para la paralelización. El gráfico también aísla la función del código Python subyacente, lo que permite que se ejecute de manera eficiente en muchos dispositivos diferentes.

Una función decorada con @tf.function se ejecuta en dos pasos:

1. En el primer paso, TensorFlow ejecuta el código de Python para la función y compila un gráfico de cálculo, lo que retrasa la ejecución de cualquier operación de TensorFlow.
2. Posteriormente, se ejecuta el gráfico de cálculo.

El primer paso se omitirá si no es necesario crear un nuevo gráfico de cálculo. Esto mejora el rendimiento de la función, pero también significa que la función no se ejecutará como el código normal de Python (en el que se ejecuta cada línea ejecutable). Por ejemplo, modifiquemos nuestra función anterior:

@tf.function
def quick_computation(x):
  print('Only prints the first time!')
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation(x)
quick_computation(x)

Esto resulta en:

Only prints the first time!

La print() solo se ejecuta una vez durante el paso de seguimiento, que es cuando se ejecuta el código Python normal. Las próximas llamadas a la función solo ejecutan operaciones de TenforFlow desde el gráfico de cálculo (operaciones de TensorFlow).

Sin embargo, si usamos tf.print() en lugar:

@tf.function
def quick_computation_with_print(x):
  tf.print("Prints every time!")
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation_with_print(x)
quick_computation_with_print(x)

Prints every time!
Prints every time!

TensorFlow incluye tf.print() en su gráfico de cálculo, ya que es una operación de TensorFlow, no una función normal de Python.

Mejores prácticas con @tf.function

Escribir código con operaciones de TensorFlow

Como acabamos de mostrar, el gráfico de cálculo ignora algunas partes del código. Esto hace que sea difícil predecir el comportamiento de la función cuando se codifica con código Python "normal", como acabamos de ver con print(). Es mejor codificar su función con operaciones de TensorFlow cuando corresponda para evitar un comportamiento inesperado.

Por ejemplo, for y while los bucles pueden o no convertirse en el bucle TensorFlow equivalente. Por lo tanto, es mejor escribir su bucle "for" como una operación vectorizada, si es posible. Esto mejorará el rendimiento de su código y garantizará que su función se rastree correctamente.

Como ejemplo, considere lo siguiente:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

@tf.function
def function_with_for(x):
    summ = float(0)
    for row in x:
      summ = summ + tf.reduce_mean(row)
    return summ

@tf.function
def vectorized_function(x):
  result = tf.reduce_mean(x, axis=0)
  return tf.reduce_sum(result)


print(function_with_for(x))
print(vectorized_function(x))

%timeit function_with_for(x)
%timeit vectorized_function(x)

tf.Tensor(0.672811, shape=(), dtype=float32)
tf.Tensor(0.67281103, shape=(), dtype=float32)
1.58 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
440 µs ± 8.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

El código con las operaciones de TensorFlow es considerablemente más rápido.

Evite las referencias a variables globales

Considere el siguiente código:

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = 2

@tf.function
def power(x):
  return tf.pow(x,y)

print(power(x))

y = 3

print(power(x))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(4.0, shape=(), dtype=float32)

La primera vez que la función decorada power() fue llamado, el valor de salida fue el esperado 4. Sin embargo, la segunda vez, la función ignoró que el valor de y fue cambiado. Esto sucede porque el valor de las variables globales de Python se congela para la función después del seguimiento.

Una mejor forma sería utilizar tf.Variable() para todas sus variables y pase ambas como argumentos a su función.

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = tf.Variable(2, dtype = tf.dtypes.float32)

@tf.function
def power(x,y):
  return tf.pow(x,y)

print(power(x,y))

y.assign(3)

print(power(x,y))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(8.0, shape=(), dtype=float32)

Depuración GME@dhr-rgv.com_s

En general, desea depurar su función en modo entusiasta y luego decorarla con @tf.function después de que su código se ejecute correctamente porque los mensajes de error en modo entusiasta son más informativos.

Algunos problemas comunes son los errores de tipo y los errores de forma. Los errores de tipo ocurren cuando hay una discrepancia en el tipo de las variables involucradas en una operación:

x = tf.Variable(1, dtype = tf.dtypes.float32)
y = tf.Variable(1, dtype = tf.dtypes.int32)

z = tf.add(x,y)

InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a float tensor but is a int32 tensor [Op:AddV2]

Los errores de tipo aparecen fácilmente y se pueden solucionar fácilmente convirtiendo una variable en un tipo diferente:

y = tf.cast(y, tf.dtypes.float32)
z = tf.add(x, y) 
tf.print(z) 

Los errores de forma ocurren cuando sus tensores no tienen la forma que requiere su operación:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)
y = tf.random.uniform(shape=[1, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

z = tf.matmul(x,y)

InvalidArgumentError: Matrix size-incompatible: In[0]: [100,100], In[1]: [1,100] [Op:MatMul]

Una herramienta conveniente para corregir ambos tipos de errores es el depurador interactivo de Python, al que puede llamar automáticamente en un Jupyter Notebook usando %pdb. Usando eso, puede codificar su función y ejecutarla a través de algunos casos de uso comunes. Si hay un error, se abre un mensaje interactivo. Este indicador le permite subir y bajar las capas de abstracción en su código y verificar los valores, tipos y formas de sus variables de TensorFlow.

Conclusión

Hemos visto cómo TensorFlow tf.function() hace que su función sea más eficiente, y cómo el @tf.function el decorador aplica la función a la suya.

Esta aceleración es útil en funciones que se llamarán muchas veces, como pasos de entrenamiento personalizados para modelos de aprendizaje automático.