Introdução aos gráficos e tf.function

Veja no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Visão geral

Este guia vai além da superfície do TensorFlow e do Keras para demonstrar como o TensorFlow funciona. Se, em vez disso, você quiser começar imediatamente com o Keras, confira a coleção de guias do Keras .

Neste guia, você aprenderá como o TensorFlow permite fazer alterações simples em seu código para obter gráficos, como os gráficos são armazenados e representados e como você pode usá-los para acelerar seus modelos.

Esta é uma visão geral que cobre como tf.function permite que você alterne da execução antecipada para a execução do gráfico. Para uma especificação mais completa de tf.function , vá para o guia tf.function .

O que são gráficos?

Nos três guias anteriores, você executou o TensorFlow com entusiasmo . Isso significa que as operações do TensorFlow são executadas pelo Python, operação por operação, e retornam os resultados de volta ao Python.

Embora a execução rápida tenha várias vantagens exclusivas, a execução gráfica permite a portabilidade fora do Python e tende a oferecer melhor desempenho. A execução do gráfico significa que os cálculos do tensor são executados como um gráfico do TensorFlow , às vezes chamado de tf.Graph ou simplesmente "gráfico".

Gráficos 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 tf.Graph . Como esses gráficos são estruturas de dados, eles podem ser salvos, executados e restaurados sem o código Python original.

É assim que um gráfico do TensorFlow que representa uma rede neural de duas camadas se parece quando visualizado no TensorBoard.

Um gráfico simples do TensorFlow

Os benefícios dos gráficos

Com um gráfico, você tem muita flexibilidade. Você pode usar o gráfico do TensorFlow em ambientes que não têm um interpretador Python, como aplicativos móveis, dispositivos incorporados e servidores de back-end. O TensorFlow usa gráficos como formato para modelos salvos quando os exporta do Python.

Os gráficos também são facilmente otimizados, permitindo que o compilador faça transformações como:

  • Inferir estaticamente o valor dos tensores dobrando nós constantes em sua computação ("dobragem constante") .
  • Separe as subpartes de uma computação que são independentes e as divida entre threads ou dispositivos.
  • Simplifique as operações aritméticas eliminando subexpressões comuns.

Existe todo um sistema de otimização, o Grappler , para realizar esta e outras acelerações.

Resumindo, os gráficos são extremamente úteis e permitem que seu TensorFlow seja executado rapidamente , executado em paralelo e executado com eficiência em vários dispositivos .

No entanto, você ainda deseja definir seus modelos de aprendizado de máquina (ou outros cálculos) em Python por conveniência e, em seguida, construir gráficos automaticamente quando precisar deles.

Configurar

import tensorflow as tf
import timeit
from datetime import datetime

Aproveitando os gráficos

Você cria e executa um gráfico no TensorFlow usando tf.function , como uma chamada direta ou como um decorador. tf.function recebe uma função regular como entrada e retorna uma Function . Uma Function é um callable do Python que cria gráficos do TensorFlow a partir da função do Python. Você usa uma Function da mesma maneira que seu equivalente em 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 se parece com uma função normal que você escreve usando operações do TensorFlow. Por baixo , no entanto, é muito diferente . Uma Function encapsula vários tf.Graph s atrás de uma API . É assim que o Function é capaz de oferecer os benefícios da execução do gráfico , como velocidade e capacidade de implantação.

tf.function se aplica a uma função e todas as outras funções que ela 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()
array([[12.]], dtype=float32)

Se você usou o TensorFlow 1.x, notará que em nenhum momento precisou definir um Placeholder ou tf.Session .

Convertendo funções Python em gráficos

Qualquer função que você escrever com o TensorFlow conterá uma mistura de operações TF integradas e lógica Python, como cláusulas if-then , loops, break , return , continue e muito mais. Embora as operações do TensorFlow sejam facilmente capturadas por um tf.Graph , a lógica específica do Python precisa passar por uma etapa extra para se tornar parte do gráfico. tf.function usa uma biblioteca chamada AutoGraph ( tf.autograph ) para converter código Python em código gerador de gráfico.

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())
First branch, with graph: 1
Second branch, with graph: 0

Embora seja improvável que você precise visualizar gráficos diretamente, você pode inspecionar as saídas para verificar os resultados exatos. Estes não são fáceis de ler, então não há necessidade de olhar com muito cuidado!

# This is the graph-generating output of AutoGraph.
print(tf.autograph.to_code(simple_relu))
def tf__simple_relu(x):
    with ag__.FunctionScope('simple_relu', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()

        def get_state():
            return (do_return, retval_)

        def set_state(vars_):
            nonlocal retval_, do_return
            (do_return, retval_) = vars_

        def if_body():
            nonlocal retval_, do_return
            try:
                do_return = True
                retval_ = ag__.ld(x)
            except:
                do_return = False
                raise

        def else_body():
            nonlocal retval_, do_return
            try:
                do_return = True
                retval_ = 0
            except:
                do_return = False
                raise
        ag__.if_stmt(ag__.converted_call(ag__.ld(tf).greater, (ag__.ld(x), 0), None, fscope), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
        return fscope.ret(retval_, do_return)
# This is the graph itself.
print(tf_simple_relu.get_concrete_function(tf.constant(1)).graph.as_graph_def())
node {
  name: "x"
  op: "Placeholder"
  attr {
    key: "_user_specified_name"
    value {
      s: "x"
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
}
node {
  name: "Greater/y"
  op: "Const"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_INT32
        tensor_shape {
        }
        int_val: 0
      }
    }
  }
}
node {
  name: "Greater"
  op: "Greater"
  input: "x"
  input: "Greater/y"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "cond"
  op: "StatelessIf"
  input: "Greater"
  input: "x"
  attr {
    key: "Tcond"
    value {
      type: DT_BOOL
    }
  }
  attr {
    key: "Tin"
    value {
      list {
        type: DT_INT32
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_BOOL
        type: DT_INT32
      }
    }
  }
  attr {
    key: "_lower_using_switch_merge"
    value {
      b: true
    }
  }
  attr {
    key: "_read_only_resource_inputs"
    value {
      list {
      }
    }
  }
  attr {
    key: "else_branch"
    value {
      func {
        name: "cond_false_34"
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
        }
        shape {
        }
      }
    }
  }
  attr {
    key: "then_branch"
    value {
      func {
        name: "cond_true_33"
      }
    }
  }
}
node {
  name: "cond/Identity"
  op: "Identity"
  input: "cond"
  attr {
    key: "T"
    value {
      type: DT_BOOL
    }
  }
}
node {
  name: "cond/Identity_1"
  op: "Identity"
  input: "cond:1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "Identity"
  op: "Identity"
  input: "cond/Identity_1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
library {
  function {
    signature {
      name: "cond_false_34"
      input_arg {
        name: "cond_placeholder"
        type: DT_INT32
      }
      output_arg {
        name: "cond_identity"
        type: DT_BOOL
      }
      output_arg {
        name: "cond_identity_1"
        type: DT_INT32
      }
    }
    node_def {
      name: "cond/Const"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
    }
    node_def {
      name: "cond/Const_1"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
    }
    node_def {
      name: "cond/Const_2"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_INT32
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_INT32
            tensor_shape {
            }
            int_val: 0
          }
        }
      }
    }
    node_def {
      name: "cond/Const_3"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
    }
    node_def {
      name: "cond/Identity"
      op: "Identity"
      input: "cond/Const_3:output:0"
      attr {
        key: "T"
        value {
          type: DT_BOOL
        }
      }
    }
    node_def {
      name: "cond/Const_4"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_INT32
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_INT32
            tensor_shape {
            }
            int_val: 0
          }
        }
      }
    }
    node_def {
      name: "cond/Identity_1"
      op: "Identity"
      input: "cond/Const_4:output:0"
      attr {
        key: "T"
        value {
          type: DT_INT32
        }
      }
    }
    ret {
      key: "cond_identity"
      value: "cond/Identity:output:0"
    }
    ret {
      key: "cond_identity_1"
      value: "cond/Identity_1:output:0"
    }
    attr {
      key: "_construction_context"
      value {
        s: "kEagerRuntime"
      }
    }
    arg_attr {
      key: 0
      value {
        attr {
          key: "_output_shapes"
          value {
            list {
              shape {
              }
            }
          }
        }
      }
    }
  }
  function {
    signature {
      name: "cond_true_33"
      input_arg {
        name: "cond_identity_1_x"
        type: DT_INT32
      }
      output_arg {
        name: "cond_identity"
        type: DT_BOOL
      }
      output_arg {
        name: "cond_identity_1"
        type: DT_INT32
      }
    }
    node_def {
      name: "cond/Const"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_BOOL
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_BOOL
            tensor_shape {
            }
            bool_val: true
          }
        }
      }
    }
    node_def {
      name: "cond/Identity"
      op: "Identity"
      input: "cond/Const:output:0"
      attr {
        key: "T"
        value {
          type: DT_BOOL
        }
      }
    }
    node_def {
      name: "cond/Identity_1"
      op: "Identity"
      input: "cond_identity_1_x"
      attr {
        key: "T"
        value {
          type: DT_INT32
        }
      }
    }
    ret {
      key: "cond_identity"
      value: "cond/Identity:output:0"
    }
    ret {
      key: "cond_identity_1"
      value: "cond/Identity_1:output:0"
    }
    attr {
      key: "_construction_context"
      value {
        s: "kEagerRuntime"
      }
    }
    arg_attr {
      key: 0
      value {
        attr {
          key: "_output_shapes"
          value {
            list {
              shape {
              }
            }
          }
        }
      }
    }
  }
}
versions {
  producer: 898
  min_consumer: 12
}

Na maioria das vezes, tf.function funcionará sem considerações especiais. No entanto, existem algumas ressalvas, e o guia tf.function pode ajudar aqui, bem como a referência completa do AutoGraph

Polimorfismo: uma Function , muitos gráficos

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

Cada vez que você invoca uma Function com novos dtypes e formas em seus argumentos, Function cria um novo tf.Graph para os novos argumentos. Os dtypes e formas das entradas de um tf.Graph são conhecidos como assinatura de entrada ou apenas assinatura .

A Function armazena o tf.Graph correspondente a essa assinatura em um ConcreteFunction . Um ConcreteFunction é um wrapper em torno de um 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.])))
tf.Tensor(5.5, shape=(), dtype=float32)
tf.Tensor([1. 0.], shape=(2,), dtype=float32)
tf.Tensor([3. 0.], shape=(2,), dtype=float32)

Se a Function já foi chamada com essa assinatura, Function não cria 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.])`.
tf.Tensor(0.0, shape=(), dtype=float32)
tf.Tensor([0. 1.], shape=(2,), dtype=float32)

Como é apoiada por vários gráficos, uma Function é polimórfica . Isso permite que ele suporte mais tipos de entrada do que um único tf.Graph poderia representar, bem como otimizar cada tf.Graph para melhor 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())
my_relu(x)
  Args:
    x: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()

my_relu(x=[1, -1])
  Returns:
    float32 Tensor, shape=(2,)

my_relu(x)
  Args:
    x: float32 Tensor, shape=(2,)
  Returns:
    float32 Tensor, shape=(2,)

Usando tf.function

Até agora, você aprendeu como converter uma função Python em um gráfico simplesmente usando tf.function como um decorador ou wrapper. Mas, na prática, fazer o tf.function funcionar corretamente pode ser complicado! Nas seções a seguir, você aprenderá como fazer seu código funcionar conforme o esperado com tf.function .

Execução gráfica versus execução ansiosa

O código em uma Function pode ser executado tanto avidamente quanto como um gráfico. Por padrão, Function executa seu código como um gráfico:

@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)
tf.Tensor([1 0 4 4 7], shape=(5,), dtype=int32)
tf.Tensor([3 6 3 0 6], shape=(5,), dtype=int32)
get_MSE(y_true, y_pred)
<tf.Tensor: shape=(), dtype=int32, numpy=11>

Para verificar se o gráfico de sua Function está fazendo o mesmo cálculo que sua função Python equivalente, você pode fazê-lo executar avidamente com tf.config.run_functions_eagerly(True) . Essa é uma opção que desativa a capacidade da Function de criar e executar gráficos , em vez de executar o código normalmente.

tf.config.run_functions_eagerly(True)
get_MSE(y_true, y_pred)
<tf.Tensor: shape=(), dtype=int32, numpy=11>
# Don't forget to set it back when you are done.
tf.config.run_functions_eagerly(False)

No entanto, Function pode se comportar de maneira diferente em gráficos e execução antecipada. A função de print do Python é um exemplo de como esses dois modos diferem. Vamos verificar o que acontece quando você insere uma instrução print em sua função e a chama 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 o que está impresso:

error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
error = get_MSE(y_true, y_pred)
Calculating MSE!

A saída é surpreendente? get_MSE imprimiu apenas uma vez, embora tenha sido chamado três vezes.

Para explicar, a instrução print é executada quando Function executa o código original para criar o gráfico em um processo conhecido como "tracing" . O rastreamento captura as operações do TensorFlow em um gráfico e print não é capturada no gráfico. Esse gráfico é então executado para todas as três chamadas sem nunca executar o código Python novamente .

Como verificação de sanidade, vamos desativar a execução do gráfico para comparar:

# 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)
Calculating MSE!
Calculating MSE!
Calculating MSE!
tf.config.run_functions_eagerly(False)

print é um efeito colateral do Python , e há outras diferenças que você deve estar ciente ao converter uma função em uma Function . Saiba mais na seção Limitações do guia Melhor desempenho com tf.function .

Execução não estrita

A execução do gráfico executa apenas as operações necessárias para produzir os efeitos observáveis, que incluem:

  • O valor de retorno da função
  • Efeitos colaterais bem conhecidos documentados, como:

Esse comportamento geralmente é conhecido como "Execução não estrita" e difere da execução antecipada, que percorre todas as operações do programa, necessárias ou não.

Em particular, a verificação de erros de tempo de execução não conta como um efeito observável. Se uma operação for ignorada porque é desnecessária, ela não poderá gerar nenhum erro de tempo de execução.

No exemplo a seguir, a operação "desnecessária" tf.gather é ignorada durante a execução do gráfico, portanto, o erro de tempo de execução InvalidArgumentError não é gerado como seria na execução antecipada. Não confie em um erro sendo gerado durante a execução de um gráfico.

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.Tensor([0.], shape=(1,), dtype=float32)
@tf.function
def unused_return_graph(x):
  tf.gather(x, [1]) # unused
  return x

# Only needed operations are run during graph exection. The error is not raised.
print(unused_return_graph(tf.constant([0.0])))
tf.Tensor([0.], shape=(1,), dtype=float32)

práticas recomendadas tf.function

Pode levar algum tempo para se acostumar com o comportamento de Function . Para começar rapidamente, os usuários iniciantes devem brincar com as funções de decoração de brinquedos com @tf.function para obter experiência em passar de ansioso para execução gráfica.

Projetar para tf.function pode ser sua melhor aposta para escrever programas TensorFlow compatíveis com gráficos. Aqui estão algumas dicas:

Vendo a aceleração

tf.function geralmente melhora o desempenho do seu código, mas a quantidade de aceleração depende do tipo de computação que você executa. Pequenas computações podem ser dominadas pela sobrecarga de chamar um grafo. Você pode medir a diferença de desempenho assim:

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))
Eager execution: 2.5637862179974036
power_as_graph = tf.function(power)
print("Graph execution:", timeit.timeit(lambda: power_as_graph(x, 100), number=1000))
Graph execution: 0.6832536700021592

tf.function é comumente usado para acelerar loops de treinamento, e você pode aprender mais sobre isso em Escrevendo um loop de treinamento do zero com Keras.

Desempenho e trocas

Gráficos podem acelerar seu código, mas o processo de criá-los tem alguma sobrecarga. Para algumas funções, a criação do gráfico leva mais tempo do que a execução do gráfico. Esse investimento geralmente é pago rapidamente com o aumento de desempenho das execuções subsequentes, mas é importante estar ciente de que as primeiras etapas de qualquer treinamento de modelo grande podem ser mais lentas devido ao rastreamento.

Não importa o tamanho do seu modelo, você deve evitar o rastreamento com frequência. O guia tf.function discute como definir especificações de entrada e usar argumentos de tensor para evitar retraçar. Se você perceber que está obtendo um desempenho excepcionalmente ruim, é uma boa ideia verificar se está refazendo acidentalmente.

Quando é um rastreamento de Function ?

Para descobrir quando sua Function está rastreando, adicione uma instrução print ao seu código. Como regra geral, Function executará a instrução print toda vez que rastrear.

@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)))
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
tf.Tensor(11, shape=(), dtype=int32)
# 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))
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
Tracing!
tf.Tensor(11, shape=(), dtype=int32)

Novos argumentos do Python sempre acionam a criação de um novo gráfico, daí o rastreamento extra.

Próximos passos

Você pode aprender mais sobre tf.function na página de referência da API e seguindo o guia Melhor desempenho com tf.function .