Entendendo o decorador @tf.function do TensorFlow

Introdução

Melhorar o desempenho de um loop de treinamento pode economizar horas de tempo de computação ao treinar modelos de aprendizado de máquina. Uma das maneiras de melhorar o desempenho do código TensorFlow é usar o tf.function() decorador – uma alteração simples de uma linha que pode tornar suas funções muito mais rápidas.

Neste pequeno guia, explicaremos como tf.function() melhora o desempenho e dê uma olhada em algumas práticas recomendadas.

Decoradores Python e tf.função()

Em Python, um decorador é uma função que modifica o comportamento de outras funções. Por exemplo, suponha que você chame a seguinte função em uma célula de notebook:

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)

No entanto, se passarmos a função custosa para um tf.function():

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

Nós temos quicker_computation() – uma nova função que executa muito mais rápido que a anterior:

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

então, tf.function() modifica some_costly_computation() e emite o quicker_computation() função. Os decoradores também modificam as funções, então era natural fazer tf.function() também um decorador.

Usar a notação do decorador é o mesmo que chamar 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)

Como funciona tf.function() Trabalhos?

Como podemos fazer certas funções rodarem 2-3x mais rápido?

O código do TensorFlow pode ser executado em dois modos: modo ansioso e modo gráfico. O modo ansioso é a maneira padrão e interativa de executar o código: toda vez que você chama uma função, ela é executada.

O modo gráfico, no entanto, é um pouco diferente. No modo gráfico, antes de executar a função, o TensorFlow cria um gráfico de computação, que é uma estrutura de dados que contém as operações necessárias para executar a função. O gráfico de computação permite que o TensorFlow simplifique os cálculos e encontre oportunidades de paralelização. O gráfico também isola a função do código Python sobreposto, permitindo que ela seja executada com eficiência em muitos dispositivos diferentes.

Uma função decorada com @tf.function é executado em duas etapas:

1. Na primeira etapa, o TensorFlow executa o código Python para a função e compila um gráfico de computação, atrasando a execução de qualquer operação do TensorFlow.
2. Em seguida, o gráfico de computação é executado.

A primeira etapa será ignorada se não houver necessidade de criar um novo gráfico de computação. Isso melhora o desempenho da função, mas também significa que a função não será executada como o código Python normal (no qual cada linha executável é executada). Por exemplo, vamos modificar nossa função 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)

Isto resulta em:

Only prints the first time!

A print() é executado apenas uma vez durante a etapa de rastreamento, que é quando o código Python normal é executado. As próximas chamadas para a função executam apenas operações TenforFlow do gráfico de computação (operações TenforFlow).

No entanto, se usarmos tf.print() em vez de:

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

O TensorFlow inclui tf.print() em seu gráfico de computação, pois é uma operação do TensorFlow – não uma função regular do Python.

Práticas recomendadas com @tf.function

Como escrever código com operações do TensorFlow

Como acabamos de mostrar, algumas partes do código são ignoradas pelo gráfico de computação. Isso torna difícil prever o comportamento da função ao codificar com código Python “normal”, como acabamos de ver com print(). É melhor codificar sua função com operações do TensorFlow quando aplicável para evitar comportamentos inesperados.

Por exemplo, a for e while loops podem ou não ser convertidos no loop equivalente do TensorFlow. Portanto, é melhor escrever seu loop “for” como uma operação vetorizada, se possível. Isso melhorará o desempenho do seu código e garantirá que sua função seja rastreada corretamente.

Como exemplo, considere o seguinte:

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)

O código com as operações do TensorFlow é consideravelmente mais rápido.

Evite Referências a Variáveis ​​Globais

Considere o seguinte 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)

A primeira vez que a função decorada power() foi chamado, o valor de saída foi o esperado 4. No entanto, na segunda vez, a função ignorou que o valor de y foi alterado. Isso acontece porque o valor das variáveis ​​globais do Python é congelado para a função após o rastreamento.

A melhor maneira seria usar tf.Variable() para todas as suas variáveis ​​e passe ambos como argumentos para sua função.

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)

depuração [email protegido]_s

Em geral, você deseja depurar sua função no modo ansioso e depois decorá-la com @tf.function depois que seu código estiver sendo executado corretamente porque as mensagens de erro no modo ansioso são mais informativas.

Alguns problemas comuns são erros de tipo e erros de forma. Erros de tipo acontecem quando há uma incompatibilidade no tipo das variáveis ​​envolvidas em uma operação:

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]

Erros de tipo aparecem facilmente e podem ser facilmente corrigidos lançando uma variável para um tipo diferente:

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

Erros de forma acontecem quando seus tensores não têm a forma que sua operação requer:

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]

Uma ferramenta conveniente para corrigir os dois tipos de erros é o depurador interativo do Python, que você pode chamar automaticamente em um Jupyter Notebook usando %pdb. Usando isso, você pode codificar sua função e executá-la em alguns casos de uso comuns. Se houver um erro, um prompt interativo será aberto. Esse prompt permite que você suba e desça as camadas de abstração em seu código e verifique os valores, tipos e formas de suas variáveis ​​do TensorFlow.

Conclusão

Vimos como o TensorFlow tf.function() torna sua função mais eficiente, e como o @tf.function decorador aplica a função ao seu próprio.

Essa aceleração é útil em funções que serão chamadas muitas vezes, como etapas de treinamento personalizadas para modelos de aprendizado de máquina.