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

Otimização de grafo do TensorFlow com o Grappler

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

O TensorFlow usa execuções no modo grafo e adiantado (eager) para fazer as computações. Um tf.Graph contém um conjunto de objetos tf.Operation (operações), 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 padrão de otimização de grafo do runtime do TensorFlow e aplica otimizações no modo grafo (dentro de tf.function) para aumentar o desempenho das computações do TensorFlow por meio de simplificações do grafo e outras otimizações de alto nível, como embutir o corpo de funções para ativar otimizações interprocedurais. Otimizar o tf.Graph também reduz o pico de uso de memória e melhora a utilização do hardware ao otimizar o mapeamento dos nós do grafo para os recursos computacionais.

Use tf.config.optimizer.set_experimental_options() para ter um controle mais granular das otimizações do tf.Graph.

Otimizadores de grafo disponíveis

O Grappler faz as otimizações de grafo por meio de um driver chamado MetaOptimizer. Os otimizadores de grafo abaixo estão disponíveis no TensorFlow:

• Otimizador com constant folding – infere estatisticamente o valor dos tensores quando possível fazendo o folding de nós constantes no grafo e materializa o resultado usando constantes.

• Otimizador aritmético – simplifica as operações aritméticas eliminando subexpressões comuns e simplificando declarações aritméticas.

• Otimizador de layout – otimiza layouts de tensores para executar operações que dependem do formato dos dados, como convoluções, de maneira mais eficaz.

• Otimizador de remapeamento – remapeia subgrafos para implementações mais eficientes por meio da substituição de subgrafos com ocorrência comum por kernels monolíticos combinados e otimizados.

• Otimizador de memória – analisa o grafo para avaliar o pico de uso de memória de cada operação e insere as operações de cópia de memória CPU-GPU para trocar memória de GPU por CPU a fim de reduzir o pico de uso de memória.

• Otimizador de dependência – remove ou reorganiza as dependências de controle para encurtar o caminho crítico de um passo do modelo ou permitir outras otimizações. Além disso, remove nós que não são realmente operações, como Identity.

• Otimizador de pruning – remove do grafo nós que não têm efeito na saída. Geralmente, é executado primeiro para reduzir o tamanho do grafo e acelerar o processamento de outros passos do Grappler.

• Otimizador de função – otimiza a biblioteca de funções de um programa do TensorFlow e o embute corpo de funções para ativar outras otimizações interprocedurais.

• Otimizador de formato – otimiza subgrafos que operam o formato e informações relacionadas ao formato.

• Otimizador de autoparalelização – paraleliza automaticamente os grafos por meio da divisão na dimensão de lote. Esse otimizador é desativado por padrão.

• Otimizador de loop – otimiza o fluxo de controle do grafo retirando dos loops os subgrafos que não variam nos loops e removendo operações de pilhas redundantes. Além disso, otimiza loops com contagens de execução conhecidas estatisticamente e remove ramificações mortas conhecidas estatisticamente em condicionais.

• Otimizador de alocador com escopo – introduz alocadores com escopo para reduzir o movimento de dados e consolidar algumas operações.

• Otimizador pin to host – troca pequenas operações para a CPU. Esse otimizador é desativado por padrão.

• Otimizador de precisão mista automático – converte tipos de dados em float16 quando aplicável para melhorar o desempenho. Atualmente, aplica-se somente a GPUs.

• Eliminador de depuração – elimina do grafo nós relacionados a operações de depuração, como tf.debugging.Assert, tf.debugging.check_numerics e tf.print. Esse otimizador é desativado por padrão.

Configuração

import numpy as np
import timeit
import traceback
import contextlib


import tensorflow as tf

Crie um gerenciador de contexto para ativar e desativar estados de otimizadores com facilidade.

@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 da execução com e sem o Grappler

O TensorFlow 2 e posteriores fazer execução adiantada (eager) por padrão. Use tf.function para mudar a execução padrão para o modo grafo. O Grappler executa automaticamente em segundo plano para aplicar as otimizações de grafo acima e melhorar o desempenho da execução.

Otimizador com constant folding

Vejamos um exemplo preliminar. Considere uma função que realize operações em constantes e retorne 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

Desative o otimizador com constant folding 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")

Ative o otimizador com constant folding e execute a função novamente. Observe a aceleração ao executar a 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")

Otimizador de eliminação de depuração

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

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 de eliminação 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)

tf.debugging.check_numerics gera um erro de argumento inválido por causa do argumento Inf em test_func.

Ative o otimizador de eliminação 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 de eliminação de depuração elimina o nó tf.debug.check_numerics do grafo e executa a função sem gerar nenhum erro.

Resumo

O runtime do TensorFlow usa o Grappler para otimizar os grafos automaticamente antes da execução. Use tf.config.optimizer.set_experimental_options para ativar ou desativar os diversos otimizadores de grafo.

Confira mais informações sobre o Grappler em Otimizações de grafo do TensorFlow.