Optimisation des graphes TensorFlow avec Grappler

Voir sur TensorFlow.org Exécuter dans Google Colab Voir la source sur GitHub Télécharger le cahier

Aperçu

TensorFlow utilise à la fois des exécutions graphiques et impatientes pour exécuter des calculs. Un tf.Graph contient un ensemble d'objets tf.Operation (ops) qui représentent des unités de calcul et d'objets tf.Tensor qui représentent les unités de données qui circulent entre les ops.

Grappler est le système d'optimisation de graphe par défaut dans l'environnement d'exécution TensorFlow. Grappler applique des optimisations en mode graphique (dans tf.function ) pour améliorer les performances de vos calculs TensorFlow grâce à des simplifications de graphique et à d'autres optimisations de haut niveau telles que l'intégration de corps de fonction pour permettre des optimisations inter-procédurales. L'optimisation de tf.Graph réduit également l'utilisation maximale de la mémoire de l'appareil et améliore l'utilisation du matériel en optimisant le mappage des nœuds de graphe pour calculer les ressources.

Utilisez tf.config.optimizer.set_experimental_options() pour un contrôle plus précis de vos optimisations tf.Graph .

Optimiseurs de graphiques disponibles

Grappler effectue des optimisations de graphes via un pilote de niveau supérieur appelé MetaOptimizer . Les optimiseurs de graphiques suivants sont disponibles avec TensorFlow :

  • Optimiseur de pliage constant - Déduit statiquement la valeur des tenseurs lorsque cela est possible en pliant des nœuds constants dans le graphique et matérialise le résultat à l'aide de constantes.
  • Optimiseur arithmétique - Simplifie les opérations arithmétiques en éliminant les sous-expressions courantes et en simplifiant les instructions arithmétiques.
  • Optimiseur de mise en page - Optimise les mises en page des tenseurs pour exécuter plus efficacement les opérations dépendant du format de données telles que les convolutions.
  • Optimiseur de remappage - Remappe les sous-graphes sur des implémentations plus efficaces en remplaçant les sous-graphes courants par des noyaux monolithiques fusionnés optimisés.
  • Optimiseur de mémoire - Analyse le graphique pour inspecter l'utilisation maximale de la mémoire pour chaque opération et insère des opérations de copie de mémoire CPU-GPU pour échanger la mémoire GPU vers le CPU afin de réduire l'utilisation maximale de la mémoire.
  • Optimiseur de dépendance - Supprime ou réorganise les dépendances de contrôle pour raccourcir le chemin critique d'une étape de modèle ou permet d'autres optimisations. Supprime également les nœuds qui sont effectivement no-ops tels que Identity.
  • Optimiseur d'élagage - Élague les nœuds qui n'ont aucun effet sur la sortie du graphique. Il est généralement exécuté en premier pour réduire la taille du graphique et accélérer le traitement dans les autres passes de Grappler.
  • Optimiseur de fonctions : optimise la bibliothèque de fonctions d'un programme TensorFlow et inline les corps de fonctions pour permettre d'autres optimisations inter-procédurales.
  • Optimiseur de forme - Optimise les sous-graphes qui fonctionnent sur la forme et les informations liées à la forme.
  • Optimiseur Autoparallel - Parallélise automatiquement les graphiques en les divisant le long de la dimension du lot. Cet optimiseur est désactivé par défaut.
  • Optimiseur de boucle - Optimise le flux de contrôle du graphe en extrayant les sous-graphes invariants de boucle des boucles et en supprimant les opérations de pile redondantes dans les boucles. Optimise également les boucles avec des nombres de trajets statiquement connus et supprime les branches mortes statiquement connues dans les conditionnels.
  • Optimiseur d'allocateur étendu - Introduit des allocateurs étendus pour réduire le déplacement des données et consolider certaines opérations.
  • Épingler à l'optimiseur d'hôte - Échange de petites opérations sur le CPU. Cet optimiseur est désactivé par défaut.
  • Optimiseur de précision mixte automatique - Convertit les types de données en float16, le cas échéant, pour améliorer les performances. S'applique actuellement uniquement aux GPU.
  • Debug stripper - Supprime les nœuds liés aux opérations de débogage telles que tf.debugging.Assert , tf.debugging.check_numerics et tf.print du graphique. Cet optimiseur est désactivé par défaut.

Installer

import numpy as np
import timeit
import traceback
import contextlib


import tensorflow as tf

Créez un gestionnaire de contexte pour basculer facilement entre les états de l'optimiseur.

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

Comparez les performances d'exécution avec et sans Grappler

TensorFlow 2 et au-delà s'exécute avec impatience par défaut. Utilisez tf.function pour basculer l'exécution par défaut en mode Graph. Grappler s'exécute automatiquement en arrière-plan pour appliquer les optimisations graphiques ci-dessus et améliorer les performances d'exécution.

Optimiseur de pliage constant

Comme exemple préliminaire, considérons une fonction qui effectue des opérations sur des constantes et renvoie une sortie.

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

Désactivez l'optimiseur de pliage constant et exécutez la fonction :

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

Activez l'optimiseur de pliage constant et exécutez à nouveau la fonction pour observer une accélération de l'exécution de la fonction.

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

Optimiseur de débogage

Considérez une fonction simple qui vérifie la valeur numérique de son argument d'entrée et la renvoie.

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

Tout d'abord, exécutez la fonction avec l'optimiseur de débogage désactivé.

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 une erreur d'argument non valide en raison de l'argument Inf de test_func .

Activez l'optimiseur de débogage et exécutez à nouveau la fonction.

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)

L'optimiseur de suppression de débogage supprime le nœud tf.debug.check_numerics du graphique et exécute la fonction sans générer d'erreur.

Résumé

L'environnement d'exécution TensorFlow utilise Grappler pour optimiser automatiquement les graphiques avant l'exécution. Utilisez tf.config.optimizer.set_experimental_options pour activer ou désactiver les différents optimiseurs de graphique.

Pour plus d'informations sur Grappler, consultez Optimisations des graphes TensorFlow .