Otimização de gráfico TensorFlow com Grappler

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

Visão geral

O TensorFlow usa execuções gráficas e ansiosas para executar cálculos. Um tf.Graph contém um conjunto de objetos tf.Operation (ops) que representam unidades de computação e objetos tf.Tensor que representam as unidades de dados que fluem entre as operações.

O Grappler é o sistema de otimização de gráficos padrão no tempo de execução do TensorFlow. O Grappler aplica otimizações no modo de gráfico (dentro de tf.function ) para melhorar o desempenho de seus cálculos do TensorFlow por meio de simplificações de gráficos e outras otimizações de alto nível, como corpos de funções embutidos para permitir otimizações entre procedimentos. A otimização do tf.Graph também reduz o uso de pico de memória do dispositivo e melhora a utilização do hardware otimizando o mapeamento de nós de gráfico para recursos de computação.

Use tf.config.optimizer.set_experimental_options() para um controle mais preciso sobre suas otimizações tf.Graph .

Otimizadores de gráficos disponíveis

O Grappler realiza otimizações gráficas por meio de um driver de nível superior chamado MetaOptimizer . Os seguintes otimizadores de gráfico estão disponíveis com o TensorFlow:

  • Otimizador de dobra constante - Infere estaticamente o valor dos tensores quando possível dobrando nós constantes no gráfico e materializando o resultado usando constantes.
  • Otimizador aritmético - Simplifica as operações aritméticas eliminando subexpressões comuns e simplificando as instruções aritméticas.
  • Otimizador de layout - Otimiza layouts de tensor para executar operações dependentes de formato de dados, como convoluções com mais eficiência.
  • Otimizador de remapeamento - Remapeia subgrafos para implementações mais eficientes, substituindo subgrafos que ocorrem com frequência por kernels monolíticos fundidos otimizados.
  • Otimizador de memória - analisa o gráfico para inspecionar o uso máximo de memória para cada operação e insere operações de cópia de memória CPU-GPU para trocar a memória GPU para CPU para reduzir o uso máximo de memória.
  • Otimizador de dependência - remove ou reorganiza as dependências de controle para encurtar o caminho crítico para uma etapa do modelo ou permite outras otimizações. Também remove nós que são efetivamente no-ops, como Identity.
  • Otimizador de poda - poda os nós que não têm efeito na saída do gráfico. Geralmente é executado primeiro para reduzir o tamanho do gráfico e acelerar o processamento em outras passagens do Grappler.
  • Otimizador de funções - Otimiza a biblioteca de funções de um programa TensorFlow e inline corpos de funções para permitir outras otimizações entre procedimentos.
  • Otimizador de forma - Otimiza os subgráficos que operam em informações relacionadas à forma e à forma.
  • Otimizador Autoparallel - Paraleliza automaticamente os gráficos dividindo ao longo da dimensão do lote. Este otimizador está desligado por padrão.
  • Otimizador de loop - Otimiza o fluxo de controle de gráfico ao içar subgráficos invariantes de loop para fora dos loops e removendo operações de pilha redundantes em loops. Também otimiza loops com contagens de disparo conhecidas estaticamente e remove ramificações mortas estaticamente conhecidas em condicionais.
  • Otimizador de alocador com escopo - apresenta alocadores com escopo para reduzir a movimentação de dados e consolidar algumas operações.
  • Pin to host otimizador - Troca pequenas operações na CPU. Este otimizador está desligado por padrão.
  • Otimizador de precisão misto automático - Converte tipos de dados em float16 quando aplicável para melhorar o desempenho. Atualmente aplica-se apenas a GPUs.
  • Debug stripper - Remove nós relacionados a operações de depuração como tf.debugging.Assert , tf.debugging.check_numerics e tf.print do gráfico. Este otimizador está desligado por padrão.

Configurar

import numpy as np
import timeit
import traceback
import contextlib


import tensorflow as tf

Crie um gerenciador de contexto para alternar facilmente os estados do otimizador.

@contextlib.contextmanager
def options(options):
  old_opts = tf.config.optimizer.get_experimental_options()
  tf.config.optimizer.set_experimental_options(options)
  try:
    yield
  finally:
    tf.config.optimizer.set_experimental_options(old_opts)

Compare o desempenho de execução com e sem Grappler

O TensorFlow 2 e versões posteriores são executados avidamente por padrão. Use tf.function para alternar a execução padrão para o modo Graph. O Grappler é executado automaticamente em segundo plano para aplicar as otimizações de gráfico acima e melhorar o desempenho da execução.

Otimizador de dobra constante

Como exemplo preliminar, considere uma função que realiza operações em constantes e retorna uma saída.

def test_function_1():
  @tf.function
  def simple_function(input_arg):
    print('Tracing!')
    a = tf.constant(np.random.randn(2000,2000), dtype = tf.float32)
    c = a
    for n in range(50):
      c = c@a
    return tf.reduce_mean(c+input_arg)

  return simple_function

Desligue o otimizador de dobra constante e execute a função:

with options({'constant_folding': False}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Vanilla execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")
{'constant_folding': False, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Vanilla execution: 0.0018392090000816097 s

Habilite o otimizador de dobra constante e execute a função novamente para observar uma aceleração na execução da função.

with options({'constant_folding': True}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Constant folded execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")
{'constant_folding': True, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Constant folded execution: 0.0006749789999958011 s

Otimizador do stripper de depuração

Considere uma função simples que verifica o valor numérico de seu argumento de entrada e o retorna.

def test_function_2():
  @tf.function
  def simple_func(input_arg):
    output = input_arg
    tf.debugging.check_numerics(output, "Bad!")
    return output
  return simple_func

Primeiro, execute a função com o otimizador do stripper de depuração desativado.

test_func = test_function_2()
p1 = tf.constant(float('inf'))
try:
  test_func(p1)
except tf.errors.InvalidArgumentError as e:
  traceback.print_exc(limit=2)
2021-09-22 20:34:55.871238: E tensorflow/core/kernels/check_numerics_op.cc:292] abnormal_detected_host @0x7f4878e00100 = {0, 1} Bad!
Traceback (most recent call last):
  File "/tmp/ipykernel_22954/3616845043.py", line 4, in <module>
    test_func(p1)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 885, in __call__
    result = self._call(*args, **kwds)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  Bad! : Tensor had Inf values
     [[node CheckNumerics (defined at tmp/ipykernel_22954/2241890286.py:5) ]] [Op:__inference_simple_func_131]

Errors may have originated from an input operation.
Input Source operations connected to node CheckNumerics:
 input_arg (defined at tmp/ipykernel_22954/3616845043.py:4)

Function call stack:
simple_func

tf.debugging.check_numerics gera um erro de argumento inválido devido ao argumento Inf para test_func .

Habilite o otimizador do stripper de depuração e execute a função novamente.

with options({'debug_stripper': True}):
  test_func2 = test_function_2()
  p1 = tf.constant(float('inf'))
  try:
    test_func2(p1)
  except tf.errors.InvalidArgumentError as e:
    traceback.print_exc(limit=2)

O otimizador do depurador remove o nó tf.debug.check_numerics do gráfico e executa a função sem gerar nenhum erro.

Resumo

O runtime do TensorFlow usa o Grappler para otimizar gráficos automaticamente antes da execução. Use tf.config.optimizer.set_experimental_options para habilitar ou desabilitar os vários otimizadores de gráfico.

Para obter mais informações sobre o Grappler, consulte Otimizações do TensorFlow Graph .