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

Introdução aos grafos e a tf.function

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

Este guia se aprofunda no TensorFlow e Keras para demonstrar como o TensorFlow funciona. Se, em vez disso, você quiser começar pelo Keras imediatamente, confira a coleção de guias do Keras.

Neste guia, você verá com o TensorFlow permite fazer alterações simples no código para obter grafos, como os grafos são armazenados e representados, além de como usá-los para acelerar os seus modelos.

Observação: para quem já conhece bem o TensorFlow 1.x, este guia demonstra um conjunto muito diferente de grafos.

**Esta é uma visão geral que demonstra como tf.function permite mudar da execução adiantada (eager) para a execução de grafo. **Para ver uma especificação mais completa de tf.function, confira o guia Desempenho melhor com tf.function.

O que são grafos?

Nos três guias anteriores, você executou o TensorFlow no modo adiantado (eager). Portanto, as operações do TensorFlow eram executadas pelo Python, operação por operação, retornando os resultados para o Python.

Embora a execução adiantada (eager) tenha diversas vantagens exclusivas, a execução de grafo permite portabilidade para fora do Python e costuma ter melhor desempenho. Execução de grafo significa que as computações de tensores são executadas como um grafo do TensorFlow, às vezes chamado de tf.Graph ou simplesmente "grafo".

Os grafos são estruturas de dados que contêm um conjunto de objetos tf.Operation, que representam unidades de computação; e objetos tf.Tensor, que representam as unidades de dados que fluem entre as operações. Eles são definidos em um contexto de tf.Graph. Como esses grafos são estruturas de dados, podem ser salvos, executados e restaurados sem o código Python original.

Veja abaixo um grafo do TensorFlow representando uma rede neural de duas camadas quando visualizado no TensorBoard:

Benefícios dos grafos

Com um grafo, você tem muita flexibilidade. É possível usar o grafo do TensorFlow em ambientes que não tenham um interpretador Python, como aplicativos móveis, dispositivos embarcados e servidores de back-end. O TensorFlow usa grafos como o formato de modelos salvos ao exportá-los do Python.

Também é fácil otimizar grafos, o que permite ao compilador fazer as seguintes tarefas:

• Inferir estatisticamente o valor de tensores fazendo o "constant folding" de nós em sua computação.

• Separar subpartes de uma computação que são independentes e dividi-las entre os threads ou dispositivos.

• Simplificar operações aritméticas eliminando subexpressões comuns.

Existe um sistema de otimização completo, chamado Grappler, que faz essas três tarefas e outras acelerações.

Resumindo, os grafos são muito úteis e permitem que o TensorFlow seja executado com rapidez, em paralelo e com eficiência em vários dispositivos.

Entretanto, ainda é importante definir seus modelos de aprendizado de máquina (ou outras computações) no Python por conveniência e depois construir os grafos automaticamente quando você precisar deles.

Configuração

Importe algumas bibliotecas necessárias:

import tensorflow as tf
import timeit
from datetime import datetime

Como aproveitar os grafos

Para criar e executar um grafo no TensorFlow, utilizamos tf.function, seja como uma chamada direta ou um decorador. tf.function recebe uma função comum como entrada e retorna uma Function. Uma Function (Função) é um callable que cria grafos do TensorFlow usando a função do Python. Uma Function é usada da mesma forma que seu equivalente do Python.

# Define a Python function.
def a_regular_function(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# `a_function_that_uses_a_graph` is a TensorFlow `Function`.
a_function_that_uses_a_graph = tf.function(a_regular_function)

# Make some tensors.
x1 = tf.constant([[1.0, 2.0]])
y1 = tf.constant([[2.0], [3.0]])
b1 = tf.constant(4.0)

orig_value = a_regular_function(x1, y1, b1).numpy()
# Call a `Function` like a Python function.
tf_function_value = a_function_that_uses_a_graph(x1, y1, b1).numpy()
assert(orig_value == tf_function_value)

Por fora, uma Function parece uma função comum escrita usando operações do TensorFlow. Porém, por dentro, é bem diferente. Uma Function encapsula váriostf.Graphs por trás de uma API (saiba mais na seção Polimorfismo). É assim que uma Function consegue proporcionar os benefícios da execução de grafo, como velocidade e capacidade de implantação (confira a seção Benefícios dos grafos acima).

tf.function aplica-se a uma função e a todas as outras funções que chama:

def inner_function(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# Use the decorator to make `outer_function` a `Function`.
@tf.function
def outer_function(x):
  y = tf.constant([[2.0], [3.0]])
  b = tf.constant(4.0)

  return inner_function(x, y, b)

# Note that the callable will create a graph that
# includes `inner_function` as well as `outer_function`.
outer_function(tf.constant([[1.0, 2.0]])).numpy()

Se você já tiver usado o TensorFlow 1.x, observará que em momento nenhum precisou definir um Placeholder ou tf.Session.

Como converter funções do Python em grafos

Toda função que você escreve com o TensorFlow conterá uma combinação de operações integradas do TF e lógica do Python, como declarações if-then, loops, break, return, continue e muito mais. Embora as operações do TensorFlow sejam facilmente capturadas por um tf.Graph, lógicas específicas do Python precisam passar por um passo extra para se tornarem parte do grafo. tf.function usa uma biblioteca chamada AutoGraph (tf.autograph) para converter código Python em código gerador de grafo.

def simple_relu(x):
  if tf.greater(x, 0):
    return x
  else:
    return 0

# `tf_simple_relu` is a TensorFlow `Function` that wraps `simple_relu`.
tf_simple_relu = tf.function(simple_relu)

print("First branch, with graph:", tf_simple_relu(tf.constant(1)).numpy())
print("Second branch, with graph:", tf_simple_relu(tf.constant(-1)).numpy())

Embora seja improvável que você precise ver grafos diretamente, pode analisar as saídas para verificar os resultados exatos, que não são fáceis de ler, então não precisa analisar detalhadamente.

# This is the graph-generating output of AutoGraph.
print(tf.autograph.to_code(simple_relu))

# This is the graph itself.
print(tf_simple_relu.get_concrete_function(tf.constant(1)).graph.as_graph_def())

Na maioria do tempo, tf.function funcionará sem considerações especiais. Porém, há algumas ressalvas, e o guia sobre tf.function pode ajudar, bem como a referência completa do AutoGraph.

Polimorfismo: uma Function, vários grafos

Um tf.Graph é especializado para um tipo de entradas (por exemplo, tensores com um dtype – tipo de dados – específico ou objetos com o mesmo id()).

Cada vez que você chama uma Function com um conjunto de argumentos que não podem ser tratados por nenhum dos grafos existentes (como argumentos com novos dtypes ou formatos incompatíveis), Function cria um novo tf.Graph especializado para esses novos argumentos. A especificação de tipo das entradas de um tf.Graph é conhecida como assinatura da entrada ou simplesmente assinatura. Para saber mais sobre quando um novo tf.Graph é gerado e como pode ser controlado, confira a seção Regras de tracing do guia Desempenho melhor com tf.function.

A Function armazena o tf.Graph que corresponde à assinatura em uma ConcreteFunction. Uma ConcreteFunction é um encapsulador de tf.Graph.

@tf.function
def my_relu(x):
  return tf.maximum(0., x)

# `my_relu` creates new graphs as it observes more signatures.
print(my_relu(tf.constant(5.5)))
print(my_relu([1, -1]))
print(my_relu(tf.constant([3., -3.])))

Se a Function já tiver sido chamada com essa assinatura, a Function não criará um novo tf.Graph.

# These two calls do *not* create new graphs.
print(my_relu(tf.constant(-2.5))) # Signature matches `tf.constant(5.5)`.
print(my_relu(tf.constant([-1., 1.]))) # Signature matches `tf.constant([3., -3.])`.

Como uma Function tem muitos grafos por trás, ela é polimorfa, o que permite receber mais tipos de entrada do que um único tf.Graph poderia representar e permite também otimizar cada tf.Graph a fim de melhorar o desempenho.

# There are three `ConcreteFunction`s (one for each graph) in `my_relu`.
# The `ConcreteFunction` also knows the return type and shape!
print(my_relu.pretty_printed_concrete_signatures())

Usando tf.function

Até o momento, você aprendeu a converter uma função do Python em um grafo apenas usando tf.function como decorador ou encapsulador. Porém, na prática, pode ser complicado fazer tf.function funcionar corretamente. Nas próximas seções, você verá como fazer seu código funcionar conforme o esperado ao usar tf.function.

Execução de grafo versus execução adiantada (eager)

O código em uma Function pode ser executado tanto no modo grafo quanto adiantado. Por padrão, Function executa o código como um grafo:

@tf.function
def get_MSE(y_true, y_pred):
  sq_diff = tf.pow(y_true - y_pred, 2)
  return tf.reduce_mean(sq_diff)

y_true = tf.random.uniform([5], maxval=10, dtype=tf.int32)
y_pred = tf.random.uniform([5], maxval=10, dtype=tf.int32)
print(y_true)
print(y_pred)

get_MSE(y_true, y_pred)

Para verificar que o grafo da sua Function esteja fazendo a mesma computação que a função equivalente do Python, você pode fazer a execução adiantada usando tf.config.run_functions_eagerly(True), que é um controle que desativa a capacidade de Functioncriar e executar grafos, que passa então a executar o código normalmente.

tf.config.run_functions_eagerly(True)

get_MSE(y_true, y_pred)

# Don't forget to set it back when you are done.
tf.config.run_functions_eagerly(False)

Porém, Function pode se comportar de maneira diferente no modo grafo e no modo adiantado (eager). A função print do Python é um exemplo de como esses dois modos diferem. Vamos conferir o que acontece ao inserir uma declaração print em sua função e chamá-la repetidamente.

@tf.function
def get_MSE(y_true, y_pred):
  print("Calculating MSE!")
  sq_diff = tf.pow(y_true - y_pred, 2)
  return tf.reduce_mean(sq_diff)

Observe a saída:

error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)

A saída é surpreendente? get_MSE só foi exibido uma vez via print, embora tenha sido chamado três vezes.

A explicação é que a declaração print é executada quando a Function executa o código original para criar o grafo em um processo conhecido como "tracing" (confira a seção Tracing do guia sobre tf.function). O tracing captura as operações do TensorFlow em um grafo, e print não é capturado no grafo. Então, esse grafo é executado para todas as três chamadas sem executar o código Python novamente.

Para testar, vamos desativar a execução de grafo a fim de fazer uma comparação:

# Now, globally set everything to run eagerly to force eager execution.
tf.config.run_functions_eagerly(True)

# Observe what is printed below.
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)

tf.config.run_functions_eagerly(False)

print é um efeito colateral do Python, e você precisa estar ciente de outras diferenças ao converter uma função em uma Function. Saiba mais na seção Limitações do guia Desempenho melhor com tf.function.

Observação: se você quiser exibir valores via print tanto na execução adiantada (eager) quanto de grafo, utilize tf.print.

Execução não estrita

O modo grafo executa somente as operações necessárias para gerar os efeitos observáveis, o que inclui:

• O valor de retorno da função.

• Efeitos colaterais conhecidos e documentados, como:

  • Operações de entrada/saída, como tf.print.

  • Operações de depuração, como as funções de asserção em tf.debugging.

  • Mutações de tf.Variable.

Esse comportamento é conhecido como "Execução não estrita", e difere da execução adiantada (eager), que passa por todas as operações do programa, sejam necessárias ou não.

Especificamente, a verificação de erro do runtime não conta como um efeito observável. Se uma operação for ignorada porque é desnecessária, não pode gerar erros de runtime.

No exemplo abaixo, a operação "desnecessária" tf.gather é ignorada durante a execução de grafo, então o erro do runtime InvalidArgumentError não é gerado como seria na execução adiantada (eager). Não conte com a geração de erros ao executar um grafo.

def unused_return_eager(x):
  # Get index 1 will fail when `len(x) == 1`
  tf.gather(x, [1]) # unused 
  return x

try:
  print(unused_return_eager(tf.constant([0.0])))
except tf.errors.InvalidArgumentError as e:
  # All operations are run during eager execution so an error is raised.
  print(f'{type(e).__name__}: {e}')

@tf.function
def unused_return_graph(x):
  tf.gather(x, [1]) # unused
  return x

# Only needed operations are run during graph execution. The error is not raised.
print(unused_return_graph(tf.constant([0.0])))

Práticas recomendadas ao usar tf.function

Pode demorar um tempo para se acostumar com o comportamento de Function. Para começar a usá-la rapidamente, usuários iniciantes podem decorar funções usando @tf.function para adquirirem experiência com a mudança de execução no modo adiantado (eager) para o modo grafo.

Utilizar tf.function poderá ser a melhor opção para escrever programas do TensorFlow compatíveis com grafos. Veja algumas dicas:

• Alterne entre execução no modo adiantado e grafo já no começo com tf.config.run_functions_eagerly para identificar se/quando os dois modos divergem.

• Crie tf.Variables fora da função do Python e modifique-as dentro dele. O mesmo vale para objetos que usam tf.Variable, como tf.keras.layers, tf.keras.Model e tf.keras.optimizers.

• Evite escrever funções que dependam das variáveis externas do Python, exceto tf.Variables e objetos do Keras. Saiba mais na seção *Dependência de variáveis globais e livres do Python * do guia sobre tf.function.

• Opte por escrever funções que recebam tensores e outros tipos do TensorFlow como entrada. Você pode passar outros tipos de objeto, mas tenha cuidado. Saiba mais na seção Dependência de objetos do Python do guia sobre tf.function.

• Inclua o máximo de computação possível usando tf.function para maximizar os ganhos de desempenho. Por exemplo, decore um passo de treinamento inteiro ou todo o loop de treinamento.

Como verificar a aceleração

Geralmente, tf.function melhora o desempenho do código, mas o nível de aceleração depende do tipo de computação executada. Pequenas computações podem acarretar sobrecargas ao chamar um grafo. Você pode mensurar a diferença de desempenho da seguinte forma:

x = tf.random.uniform(shape=[10, 10], minval=-1, maxval=2, dtype=tf.dtypes.int32)

def power(x, y):
  result = tf.eye(10, dtype=tf.dtypes.int32)
  for _ in range(y):
    result = tf.matmul(x, result)
  return result

print("Eager execution:", timeit.timeit(lambda: power(x, 100), number=1000), "seconds")

power_as_graph = tf.function(power)
print("Graph execution:", timeit.timeit(lambda: power_as_graph(x, 100), number=1000), "seconds")

tf.function é usada com frequência para acelerar os loops de treinamento. Saiba mais na seção Como acelerar o passo de treinamento com tf.function do guia do Keras Como escrever um loop de treinamento do zero.

Observação: você também pode tentar usar tf.function(jit_compile=True) para conseguir um aumento maior do desempenho, especialmente se o seu código tiver muito fluxo de controle do TensorFlow e usar vários tensores pequenos. Saiba mais na seção Execução explícita com tf.function(jit_compile=True) da Visão geral do XLA.

Desempenho e contrapartidas

Os grafos podem acelerar seu código, mas o processo de criá-los tem uma certa sobrecarga. Para algumas funções, a criação do grafo leva mais tempo do que a execução. Geralmente, esse investimento é pago rapidamente com o aumento de desempenho das execuções subsequentes, mas é importante ter em mente que os primeiros passos de treinamento de um modelo grande podem ser mais lentos devido ao tracing.

Não importa o tamanho do seu modelo, deve-se evitar fazer o tracing com frequência. O guia sobre tf.function demonstra como definir especificações de entrada e usar argumentos de tensores para evitar o retracing na seção Como controlar o retracing. Se você observar um desempenho baixo, é uma boa ideia verificar se está fazendo retracing acidentalmente

Quando uma Function faz tracing?

Para descobrir quando sua Function está fazendo tracing, adicione uma declaração print ao código. Por via de regra, uma Function executará a declaração print toda vez que fizer o tracing.

@tf.function
def a_function_with_python_side_effect(x):
  print("Tracing!") # An eager-only side effect.
  return x * x + tf.constant(2)

# This is traced the first time.
print(a_function_with_python_side_effect(tf.constant(2)))
# The second time through, you won't see the side effect.
print(a_function_with_python_side_effect(tf.constant(3)))

# This retraces each time the Python argument changes,
# as a Python argument could be an epoch count or other
# hyperparameter.
print(a_function_with_python_side_effect(2))
print(a_function_with_python_side_effect(3))

Novos argumentos do Python sempre acionam a criação de um novo grafo, e é daí que vem o tracing extra.

Próximos passos

Saiba mais sobre tf.function na página de referência da API e no guia Desempenho melhor com tf.function.