Avoir une question? Connectez-vous avec la communauté sur le forum TensorFlow Visiter le forum

De meilleures performances avec tf.function

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

Dans TensorFlow 2, l' exécution rapide est activée par défaut. L'interface utilisateur est intuitive et flexible (l'exécution d'opérations ponctuelles est beaucoup plus facile et rapide), mais cela peut se faire au détriment des performances et de la déployabilité.

Vous pouvez utiliser tf.function pour créer des graphiques à partir de vos programmes. Il s'agit d'un outil de transformation qui crée des graphiques de flux de données indépendants de Python à partir de votre code Python. Cela vous aidera à créer des modèles performants et portables, et il est nécessaire d'utiliser SavedModel .

Ce guide vous aidera à conceptualiser le fonctionnement de tf.function sous le capot, afin que vous puissiez l'utiliser efficacement.

Les principaux points à retenir et recommandations sont :

  • Déboguez en mode avide, puis décorez avec @tf.function .
  • Ne vous fiez pas aux effets secondaires de Python comme la mutation d'objet ou l'ajout de liste.
  • tf.function fonctionne mieux avec les opérations TensorFlow ; Les appels NumPy et Python sont convertis en constantes.

Installer

import tensorflow as tf

Définissez une fonction d'assistance pour illustrer les types d'erreurs que vous pourriez rencontrer :

import traceback
import contextlib

# Some helper code to demonstrate the kinds of errors you might encounter.
@contextlib.contextmanager
def assert_raises(error_class):
  try:
    yield
  except error_class as e:
    print('Caught expected exception \n  {}:'.format(error_class))
    traceback.print_exc(limit=2)
  except Exception as e:
    raise e
  else:
    raise Exception('Expected {} to be raised but no error was raised!'.format(
        error_class))

Notions de base

Usage

Une Function vous définissez (par exemple en appliquant le décorateur @tf.function ) est comme une opération principale de TensorFlow : vous pouvez l'exécuter avec impatience ; vous pouvez calculer des gradients ; etc.

@tf.function  # The decorator converts `add` into a `Function`.
def add(a, b):
  return a + b

add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]
<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[2., 2.],
       [2., 2.]], dtype=float32)>
v = tf.Variable(1.0)
with tf.GradientTape() as tape:
  result = add(v, 1.0)
tape.gradient(result, v)
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

Vous pouvez utiliser des Function s dans d'autres Function s.

@tf.function
def dense_layer(x, w, b):
  return add(tf.matmul(x, w), b)

dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))
<tf.Tensor: shape=(3, 2), dtype=float32, numpy=
array([[3., 3.],
       [3., 3.],
       [3., 3.]], dtype=float32)>

Function peuvent être plus rapides que le code avide, en particulier pour les graphes avec de nombreuses petites opérations. Mais pour les graphiques avec quelques opérations coûteuses (comme les convolutions), vous ne verrez peut-être pas beaucoup d'accélération.

import timeit
conv_layer = tf.keras.layers.Conv2D(100, 3)

@tf.function
def conv_fn(image):
  return conv_layer(image)

image = tf.zeros([1, 200, 200, 100])
# Warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Note how there's not much difference in performance for convolutions")
Eager conv: 0.23302616399996623
Function conv: 0.21780993200013654
Note how there's not much difference in performance for convolutions

Tracé

Cette section explique comment Function fonctionne sous le capot, y compris les détails d'implémentation qui peuvent changer à l'avenir . Cependant, une fois que vous comprenez pourquoi et quand le traçage se produit, il est beaucoup plus facile d'utiliser efficacement tf.function !

Qu'est-ce que le « traçage » ?

Une Function exécute votre programme dans un graphique TensorFlow . Cependant, un tf.Graph ne peut pas représenter tout ce que vous tf.Graph dans un programme TensorFlow avide. Par exemple, Python prend en charge le polymorphisme, mais tf.Graph nécessite que ses entrées aient un type de données et une dimension spécifiés. Ou vous pouvez effectuer des tâches annexes telles que la lecture d'arguments de ligne de commande, le déclenchement d'une erreur ou l'utilisation d'un objet Python plus complexe ; aucune de ces choses ne peut s'exécuter dans un tf.Graph .

Function comble cette lacune en séparant votre code en deux étapes :

1) Dans la première étape, appelée " traçage ", Function crée un nouveau tf.Graph . Le code Python s'exécute normalement, mais toutes les opérations TensorFlow (comme l'ajout de deux Tensors) sont différées : elles sont capturées par le tf.Graph et ne sont pas exécutées.

2) Dans la deuxième étape, un tf.Graph qui contient tout ce qui a été différé dans la première étape est exécuté. Cette étape est beaucoup plus rapide que l'étape de traçage.

Selon ses entrées, Function n'exécutera pas toujours la première étape lorsqu'elle est appelée. Voir « Règles de traçage » ci-dessous pour avoir une meilleure idée de la façon dont il prend cette décision. Ignorer la première étape et n'exécuter que la deuxième étape est ce qui vous donne les hautes performances de TensorFlow.

Lorsque Function décide de tracer, l'étape de traçage est immédiatement suivie de la deuxième étape, donc l'appel de Function crée et exécute à la tf.Graph le tf.Graph . Plus tard, vous verrez comment vous pouvez exécuter uniquement l'étape de traçage avec get_concrete_function .

Lorsque nous passons des arguments de différents types dans une Function , les deux étapes sont exécutées :

@tf.function
def double(a):
  print("Tracing with", a)
  return a + a

print(double(tf.constant(1)))
print()
print(double(tf.constant(1.1)))
print()
print(double(tf.constant("a")))
print()
Tracing with Tensor("a:0", shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Tracing with Tensor("a:0", shape=(), dtype=float32)
tf.Tensor(2.2, shape=(), dtype=float32)

Tracing with Tensor("a:0", shape=(), dtype=string)
tf.Tensor(b'aa', shape=(), dtype=string)

Notez que si vous appelez à plusieurs reprises une Function avec le même type d'argument, TensorFlow ignorera l'étape de traçage et réutilisera un graphique précédemment tracé, car le graphique généré serait identique.

# This doesn't print 'Tracing with ...'
print(double(tf.constant("b")))
tf.Tensor(b'bb', shape=(), dtype=string)

Vous pouvez utiliser pretty_printed_concrete_signatures() pour voir toutes les traces disponibles :

print(double.pretty_printed_concrete_signatures())
double(a)
  Args:
    a: int32 Tensor, shape=()
  Returns:
    int32 Tensor, shape=()

double(a)
  Args:
    a: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()

double(a)
  Args:
    a: string Tensor, shape=()
  Returns:
    string Tensor, shape=()

Jusqu'à présent, vous avez vu que tf.function crée une couche de répartition dynamique en cache sur la logique de traçage de graphique de TensorFlow. Pour être plus précis sur la terminologie :

  • Un tf.Graph est la représentation brute, tf.Graph du langage et portable d'un calcul TensorFlow.
  • Une ConcreteFunction enveloppe un tf.Graph .
  • Une Function gère un cache de ConcreteFunction s et sélectionne la bonne pour vos entrées.
  • tf.function enveloppe une fonction Python, renvoyant un objet Function .
  • Le traçage crée un tf.Graph et l'enveloppe dans une ConcreteFunction , également appelée trace.

Règles de traçage

Une Function détermine s'il faut réutiliser une ConcreteFunction tracée en calculant une clé de cache à partir des args et kwargs d'une entrée. Une clé de cache est une clé qui identifie une ConcreteFunction basée sur les arguments et les kwargs d'entrée de l'appel de Function , selon les règles suivantes (qui peuvent changer) :

  • La clé générée pour un tf.Tensor est sa forme et son type.
  • La clé générée pour un tf.Variable est un identifiant de variable unique.
  • La clé générée pour une primitive Python (comme int , float , str ) est sa valeur.
  • La clé générée pour imbriquée dict s, list s, tuple s, namedtuple s et attr s est tuple aplaties des clés de la feuille (voirnest.flatten ). (En raison de cet aplatissement, l'appel d'une fonction concrète avec une structure d'imbrication différente de celle utilisée lors du traçage entraînera une TypeError).
  • Pour tous les autres types Python, la clé est unique à l'objet. De cette façon, une fonction ou une méthode est tracée indépendamment pour chaque instance avec laquelle elle est appelée.

Contrôler le retraçage

Le retraçage, c'est-à-dire lorsque votre Function crée plusieurs traces, permet de garantir que TensorFlow génère des graphiques corrects pour chaque ensemble d'entrées. Cependant, le traçage est une opération coûteuse ! Si votre Function retrace un nouveau graphique pour chaque appel, vous constaterez que votre code s'exécute plus lentement que si vous n'aviez pas utilisé tf.function .

Pour contrôler le comportement de traçage, vous pouvez utiliser les techniques suivantes :

  • Spécifiez input_signature dans tf.function pour limiter le traçage.
@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)

print(next_collatz(tf.constant([1, 2])))
# You specified a 1-D tensor in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([[1, 2], [3, 4]]))

# You specified an int32 dtype in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([1.0, 2.0]))
Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([4 1], shape=(2,), dtype=int32)
Caught expected exception 
  <class 'ValueError'>:
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-14ebce7b7ee8>", line 9, in <module>
    next_collatz(tf.constant([[1, 2], [3, 4]]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None))
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-14ebce7b7ee8>", line 13, in <module>
    next_collatz(tf.constant([1.0, 2.0]))
ValueError: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor([1. 2.], shape=(2,), dtype=float32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None))
  • Spécifiez une dimension [Aucun] dans tf.TensorSpec pour permettre une flexibilité dans la réutilisation des traces.

    Étant donné que TensorFlow correspond aux tenseurs en fonction de leur forme, l'utilisation d'une dimension None comme caractère générique permettra aux Function s de réutiliser les traces pour une entrée de taille variable. Une entrée de taille variable peut se produire si vous avez des séquences de longueurs différentes, ou des images de tailles différentes pour chaque lot (voir les didacticiels Transformer et Deep Dream par exemple).

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def g(x):
  print('Tracing with', x)
  return x

# No retrace!
print(g(tf.constant([1, 2, 3])))
print(g(tf.constant([1, 2, 3, 4, 5])))
Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([1 2 3], shape=(3,), dtype=int32)
tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)
  • Transmettez les arguments Python aux Tensors pour réduire le retraçage.

    Souvent, les arguments Python sont utilisés pour contrôler les hyperparamètres et les constructions graphiques - par exemple, num_layers=10 ou training=True ou nonlinearity='relu' . Donc, si l'argument Python change, il est logique que vous deviez retracer le graphique.

    Cependant, il est possible qu'un argument Python ne soit pas utilisé pour contrôler la construction du graphe. Dans ces cas, une modification de la valeur Python peut déclencher un retraçage inutile. Prenez, par exemple, cette boucle d'entraînement, qu'AutoGraph va dérouler dynamiquement. Malgré les multiples traces, le graphe généré est en fait identique, donc le retraçage est inutile.

def train_one_step():
  pass

@tf.function
def train(num_steps):
  print("Tracing with num_steps = ", num_steps)
  tf.print("Executing with num_steps = ", num_steps)
  for _ in tf.range(num_steps):
    train_one_step()

print("Retracing occurs for different Python arguments.")
train(num_steps=10)
train(num_steps=20)

print()
print("Traces are reused for Tensor arguments.")
train(num_steps=tf.constant(10))
train(num_steps=tf.constant(20))
Retracing occurs for different Python arguments.
Tracing with num_steps =  10
Executing with num_steps =  10
Tracing with num_steps =  20
Executing with num_steps =  20

Traces are reused for Tensor arguments.
Tracing with num_steps =  Tensor("num_steps:0", shape=(), dtype=int32)
Executing with num_steps =  10
Executing with num_steps =  20

Si vous devez forcer le retraçage, créez un nouveau Function . Function objets Function séparés sont garantis pour ne pas partager de traces.

def f():
  print('Tracing!')
  tf.print('Executing')

tf.function(f)()
tf.function(f)()
Tracing!
Executing
Tracing!
Executing

Obtenir des fonctions concrètes

Chaque fois qu'une fonction est tracée, une nouvelle fonction concrète est créée. Vous pouvez obtenir directement une fonction concrète, en utilisant get_concrete_function .

print("Obtaining concrete trace")
double_strings = double.get_concrete_function(tf.constant("a"))
print("Executing traced function")
print(double_strings(tf.constant("a")))
print(double_strings(a=tf.constant("b")))
Obtaining concrete trace
Executing traced function
tf.Tensor(b'aa', shape=(), dtype=string)
tf.Tensor(b'bb', shape=(), dtype=string)
# You can also call get_concrete_function on an InputSpec
double_strings_from_inputspec = double.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.string))
print(double_strings_from_inputspec(tf.constant("c")))
Tracing with Tensor("a:0", shape=(), dtype=string)
tf.Tensor(b'cc', shape=(), dtype=string)

L'impression d'une ConcreteFunction affiche un résumé de ses arguments d'entrée (avec types) et de son type de sortie.

print(double_strings)
ConcreteFunction double(a)
  Args:
    a: string Tensor, shape=()
  Returns:
    string Tensor, shape=()

Vous pouvez également récupérer directement la signature d'une fonction concrète.

print(double_strings.structured_input_signature)
print(double_strings.structured_outputs)
((TensorSpec(shape=(), dtype=tf.string, name='a'),), {})
Tensor("Identity:0", shape=(), dtype=string)

L'utilisation d'une trace concrète avec des types incompatibles générera une erreur

with assert_raises(tf.errors.InvalidArgumentError):
  double_strings(tf.constant(1))
Caught expected exception 
  <class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-e4e2860a4364>", line 2, in <module>
    double_strings(tf.constant(1))
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute __inference_double_162 as input #0(zero-based) was expected to be a string tensor but is a int32 tensor [Op:__inference_double_162]

Vous pouvez remarquer que les arguments Python reçoivent un traitement spécial dans la signature d'entrée d'une fonction concrète. Avant TensorFlow 2.3, les arguments Python étaient simplement supprimés de la signature de la fonction concrète. À partir de TensorFlow 2.3, les arguments Python restent dans la signature, mais sont contraints de prendre la valeur définie lors du traçage.

@tf.function
def pow(a, b):
  return a ** b

square = pow.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)
print(square)
ConcreteFunction pow(a, b=2)
  Args:
    a: float32 Tensor, shape=<unknown>
  Returns:
    float32 Tensor, shape=<unknown>
assert square(tf.constant(10.0)) == 100

with assert_raises(TypeError):
  square(tf.constant(10.0), b=3)
Caught expected exception 
  <class 'TypeError'>:
Traceback (most recent call last):
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1725, in _call_impl
    cancellation_manager)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1770, in _call_with_flat_signature
    self._flat_signature_summary(), ", ".join(sorted(kwargs))))
TypeError: pow(a) got unexpected keyword arguments: b.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-d163f3d206cb>", line 4, in <module>
    square(tf.constant(10.0), b=3)
TypeError: ConcreteFunction pow(a, b) was constructed with int value 2 in b, but was called with int value 3

Obtention de graphiques

Chaque fonction concrète est un wrapper appelable autour d'un tf.Graph . Bien que la récupération de l'objet tf.Graph réel ne soit pas quelque chose que vous aurez normalement besoin de faire, vous pouvez l'obtenir facilement à partir de n'importe quelle fonction concrète.

graph = double_strings.graph
for node in graph.as_graph_def().node:
  print(f'{node.input} -> {node.name}')
[] -> a
['a', 'a'] -> add
['add'] -> Identity

Débogage

En général, le débogage du code est plus facile en mode impatient qu'à l'intérieur de tf.function . Vous devez vous assurer que votre code s'exécute sans erreur en mode impatient avant de décorer avec tf.function . Pour faciliter le processus de débogage, vous pouvez appeler tf.config.run_functions_eagerly(True) pour désactiver et réactiver globalement tf.function .

Lorsque vous recherchez des problèmes qui n'apparaissent que dans tf.function , voici quelques conseils :

  • Les anciens appels d' print Python ne s'exécutent que pendant le traçage, ce qui vous aide à savoir quand votre fonction est (re)tracée.
  • tf.print appels tf.print s'exécuteront à chaque fois et peuvent vous aider à retrouver les valeurs intermédiaires pendant l'exécution.
  • tf.debugging.enable_check_numerics est un moyen facile de localiser où NaNs et Inf sont créés.
  • pdb (le débogueur Python ) peut vous aider à comprendre ce qui se passe pendant le traçage. (Attention : pdb vous placera dans le code source transformé par AutoGraph.)

Transformations d'AutoGraph

AutoGraph est une bibliothèque qui est tf.function par défaut dans tf.function , et transforme un sous-ensemble de code avide Python en opérations TensorFlow compatibles avec les graphes. Cela inclut le flux de contrôle comme if , for , while .

Les opérations TensorFlow telles que tf.cond et tf.while_loop continuent de fonctionner, mais le flux de contrôle est souvent plus facile à écrire et à comprendre lorsqu'il est écrit en Python.

# A simple loop

@tf.function
def f(x):
  while tf.reduce_sum(x) > 1:
    tf.print(x)
    x = tf.tanh(x)
  return x

f(tf.random.uniform([5]))
[0.710546374 0.327660799 0.393230557 0.545059443 0.666661739]
[0.611019373 0.316417336 0.374141902 0.496808201 0.582779706]
[0.54484427 0.306263864 0.357609242 0.45960331 0.52468282]
[0.496646136 0.297034383 0.343106419 0.429760844 0.481306106]
[0.459475428 0.288596332 0.330247819 0.405121386 0.44728902]
[0.429656595 0.280842364 0.318743408 0.384322464 0.419668049]
[0.405034214 0.273684502 0.308370233 0.366455346 0.396650732]
[0.384248167 0.267049909 0.298953682 0.350887358 0.377079546]
[0.366391063 0.260877609 0.290354759 0.337162226 0.360168517]
[0.350830942 0.255116194 0.282461286 0.324941576 0.345362455]
[0.337112248 0.2497219 0.275181532 0.313968241 0.332256317]
[0.324896872 0.244657204 0.268439621 0.304042816 0.320546716]
[0.313927948 0.239889801 0.262172282 0.295007944 0.310001194]
[0.304006279 0.235391632 0.256326199 0.286737591 0.300438195]
[0.294974595 0.231138244 0.250856102 0.279129326 0.291713566]
[0.286706954 0.227108166 0.245723218 0.272099048 0.283711195]
[0.279101074 0.223282441 0.240894228 0.265576899 0.276336372]
[0.272072881 0.219644368 0.23634018 0.259504348 0.269510925]
[0.26555258 0.216179058 0.232035935 0.253831863 0.263169676]
[0.259481668 0.212873235 0.227959365 0.24851726 0.257257849]
[0.253810644 0.209715009 0.224091038 0.243524343 0.251728892]
[0.248497337 0.206693679 0.220413819 0.238821834 0.246543139]
[0.243505597 0.20379965 0.216912434 0.2343826 0.241666421]
[0.238804176 0.20102416 0.213573262 0.230182901 0.237069115]
[0.234365895 0.198359385 0.210384145 0.226201892 0.232725471]
[0.230167076 0.195798069 0.207334161 0.222421184 0.228612974]
[0.226186857 0.19333373 0.204413429 0.218824506 0.224711776]
[0.222406894 0.190960392 0.201613098 0.215397388 0.221004337]
[0.218810901 0.188672557 0.198925063 0.212126866 0.217475086]
[0.215384394 0.186465234 0.196342021 0.209001362 0.214110211]
[0.212114424 0.184333771 0.193857312 0.206010461 0.210897282]
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.20898944, 0.18227392, 0.19146483, 0.2031447 , 0.20782518],
      dtype=float32)>

Si vous êtes curieux, vous pouvez inspecter le code généré par l'autographe.

print(tf.autograph.to_code(f.python_function))
def tf__f(x):
    with ag__.FunctionScope('f', '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 (x,)

        def set_state(vars_):
            nonlocal x
            (x,) = vars_

        def loop_body():
            nonlocal x
            ag__.converted_call(ag__.ld(tf).print, (ag__.ld(x),), None, fscope)
            x = ag__.converted_call(ag__.ld(tf).tanh, (ag__.ld(x),), None, fscope)

        def loop_test():
            return (ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) > 1)
        ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('x',), {})
        try:
            do_return = True
            retval_ = ag__.ld(x)
        except:
            do_return = False
            raise
        return fscope.ret(retval_, do_return)

Conditionnels

AutoGraph convertira certaines instructions if <condition> en appels tf.cond équivalents. Cette substitution est effectuée si <condition> est un Tenseur. Sinon, l'instruction if est exécutée en tant que conditionnel Python.

Une conditionnelle Python s'exécute pendant le traçage, donc exactement une branche de la conditionnelle sera ajoutée au graphique. Sans AutoGraph, ce graphique tracé serait incapable de prendre la branche alternative s'il existe un flux de contrôle dépendant des données.

tf.cond trace et ajoute les deux branches du conditionnel au graphe, en sélectionnant dynamiquement une branche au moment de l'exécution. Le traçage peut avoir des effets secondaires inattendus ; consultez les effets de traçage AutoGraph pour plus d'informations.

@tf.function
def fizzbuzz(n):
  for i in tf.range(1, n + 1):
    print('Tracing for loop')
    if i % 15 == 0:
      print('Tracing fizzbuzz branch')
      tf.print('fizzbuzz')
    elif i % 3 == 0:
      print('Tracing fizz branch')
      tf.print('fizz')
    elif i % 5 == 0:
      print('Tracing buzz branch')
      tf.print('buzz')
    else:
      print('Tracing default branch')
      tf.print(i)

fizzbuzz(tf.constant(5))
fizzbuzz(tf.constant(20))
Tracing for loop
Tracing fizzbuzz branch
Tracing fizz branch
Tracing buzz branch
Tracing default branch
1
2
fizz
4
buzz
1
2
fizz
4
buzz
fizz
7
8
fizz
buzz
11
fizz
13
14
fizzbuzz
16
17
fizz
19
buzz

Consultez la documentation de référence pour des restrictions supplémentaires sur les instructions if converties par AutoGraph.

Boucles

AutoGraph convertira certaines instructions for et while en opérations de boucle TensorFlow équivalentes, comme tf.while_loop . Si non converti, le for ou while boucle est exécutée en boucle Python.

Cette substitution s'effectue dans les situations suivantes :

  • for x in y : si y est un Tensor, convertissez-le en tf.while_loop . Dans le cas particulier où y est untf.data.Dataset , une combinaison d'opérationstf.data.Dataset est générée.
  • while <condition> : si <condition> est un Tensor, convertissez-le en tf.while_loop .

Une boucle Python s'exécute pendant le traçage, ajoutant des opérations supplémentaires au tf.Graph pour chaque itération de la boucle.

Une boucle TensorFlow trace le corps de la boucle et sélectionne dynamiquement le nombre d'itérations à exécuter au moment de l'exécution. Le corps de la boucle n'apparaît qu'une seule fois dans le tf.Graph généré.

Consultez la documentation de référence pour des restrictions supplémentaires sur les instructions for et while converties for AutoGraph.

Boucle sur les données Python

Un écueil courant consiste à boucler les données Python/NumPy dans un tf.function . Cette boucle s'exécutera pendant le processus de traçage, ajoutant une copie de votre modèle au tf.Graph pour chaque itération de la boucle.

Si vous souhaitez envelopper toute la boucle d'entraînement dans tf.function , le moyen le plus sûr de le faire est d'envelopper vos données en tant quetf.data.Dataset afin qu'AutoGraph déroule dynamiquement la boucle d'entraînement.

def measure_graph_size(f, *args):
  g = f.get_concrete_function(*args).graph
  print("{}({}) contains {} nodes in its graph".format(
      f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node)))

@tf.function
def train(dataset):
  loss = tf.constant(0)
  for x, y in dataset:
    loss += tf.abs(y - x) # Some dummy computation.
  return loss

small_data = [(1, 1)] * 3
big_data = [(1, 1)] * 10
measure_graph_size(train, small_data)
measure_graph_size(train, big_data)

measure_graph_size(train, tf.data.Dataset.from_generator(
    lambda: small_data, (tf.int32, tf.int32)))
measure_graph_size(train, tf.data.Dataset.from_generator(
    lambda: big_data, (tf.int32, tf.int32)))
train([(1, 1), (1, 1), (1, 1)]) contains 11 nodes in its graph
train([(1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]) contains 32 nodes in its graph
train(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 10 nodes in its graph
train(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 10 nodes in its graph

Lors de l' tf.data.Dataset.from_generator de données Python/NumPy dans un ensemble de données, tf.data.Dataset.from_generator compte de tf.data.Dataset.from_generator par rapport à tf.data.Dataset.from_tensors . Le premier conservera les données en Python et les récupérera via tf.py_function ce qui peut avoir des implications sur les performances, tandis que le second regroupera une copie des données sous la forme d'un grand nœud tf.constant() dans le graphique, ce qui peut avoir des implications sur la mémoire.

La lecture de données à partir de fichiers via TFRecordDataset , CsvDataset , etc. est le moyen le plus efficace de consommer des données, car TensorFlow lui-même peut gérer le chargement et la prélecture asynchrone des données, sans avoir à impliquer Python. Pour en savoir plus, consultez le tf.data : Build TensorFlow input pipelines .

Accumuler des valeurs dans une boucle

Un modèle courant consiste à accumuler des valeurs intermédiaires à partir d'une boucle. Normalement, cela est accompli en ajoutant à une liste Python ou en ajoutant des entrées à un dictionnaire Python. Cependant, comme il s'agit d'effets secondaires Python, ils ne fonctionneront pas comme prévu dans une boucle déroulée dynamiquement. Utilisez tf.TensorArray pour accumuler les résultats d'une boucle déroulée dynamiquement.

batch_size = 2
seq_len = 3
feature_size = 4

def rnn_step(inp, state):
  return inp + state

@tf.function
def dynamic_rnn(rnn_step, input_data, initial_state):
  # [batch, time, features] -> [time, batch, features]
  input_data = tf.transpose(input_data, [1, 0, 2])
  max_seq_len = input_data.shape[0]

  states = tf.TensorArray(tf.float32, size=max_seq_len)
  state = initial_state
  for i in tf.range(max_seq_len):
    state = rnn_step(input_data[i], state)
    states = states.write(i, state)
  return tf.transpose(states.stack(), [1, 0, 2])

dynamic_rnn(rnn_step,
            tf.random.uniform([batch_size, seq_len, feature_size]),
            tf.zeros([batch_size, feature_size]))
<tf.Tensor: shape=(2, 3, 4), dtype=float32, numpy=
array([[[0.60458577, 0.3308612 , 0.7878152 , 0.3223114 ],
        [0.9110272 , 1.0819752 , 1.7657743 , 1.2409766 ],
        [1.7235098 , 1.5416101 , 2.2929285 , 1.9181627 ]],

       [[0.89487076, 0.22811687, 0.342862  , 0.5752872 ],
        [1.0133923 , 0.28650808, 0.9558767 , 1.0829899 ],
        [1.9280962 , 1.1437279 , 0.9857702 , 1.4834155 ]]], dtype=float32)>

Limites

TensorFlow Function a quelques limitations de par sa conception dont vous devez être conscient lors de la conversion d'une fonction Python en Function .

Exécuter les effets secondaires de Python

Les effets secondaires, tels que l'impression, l'ajout à des listes et la mutation des variables globales, peuvent se comporter de manière inattendue dans une Function , s'exécutant parfois deux fois ou pas tous. Ils ne se produisent que la première fois que vous appelez une Function avec un ensemble d'entrées. Ensuite, le tf.Graph tracé est tf.Graph , sans exécuter le code Python.

La règle générale est d'éviter de compter sur les effets secondaires Python dans votre logique et de ne les utiliser que pour déboguer vos traces. Sinon, les API TensorFlow comme tf.data , tf.print , tf.summary , tf.Variable.assign et tf.TensorArray sont le meilleur moyen de garantir que votre code sera exécuté par l'environnement d'exécution TensorFlow à chaque appel.

@tf.function
def f(x):
  print("Traced with", x)
  tf.print("Executed with", x)

f(1)
f(1)
f(2)
Traced with 1
Executed with 1
Executed with 1
Traced with 2
Executed with 2

Si vous souhaitez exécuter du code Python lors de chaque invocation d'une Function , tf.py_function est une trappe de sortie. L'inconvénient de tf.py_function est qu'il n'est pas portable ou particulièrement performant, ne peut pas être enregistré avec SavedModel et ne fonctionne pas bien dans les configurations distribuées (multi-GPU, TPU). De plus, puisque tf.py_function doit être câblé dans le graphique, il convertit toutes les entrées/sorties en tenseurs.

Modification des variables globales et libres de Python

La modification des variables globales et libres de Python compte comme un effet secondaire de Python, elle ne se produit donc que pendant le traçage.

external_list = []

@tf.function
def side_effect(x):
  print('Python side effect')
  external_list.append(x)

side_effect(1)
side_effect(1)
side_effect(1)
# The list append only happened once!
assert len(external_list) == 1
Python side effect

Vous devez éviter de muter des conteneurs tels que des listes, des dicts, d'autres objets qui vivent en dehors de la Function . Utilisez plutôt des arguments et des objets TF. Par exemple, la section « Accumuler des valeurs dans une boucle » contient un exemple de la manière dont les opérations de type liste peuvent être implémentées.

Vous pouvez, dans certains cas, capturer et manipuler l'état s'il s'agit d'un tf.Variable . C'est ainsi que les poids des modèles Keras sont mis à jour avec des appels répétés à la même ConcreteFunction .

Utiliser des itérateurs et des générateurs Python

De nombreuses fonctionnalités Python, telles que les générateurs et les itérateurs, s'appuient sur l'environnement d'exécution Python pour suivre l'état. En général, bien que ces constructions fonctionnent comme prévu en mode avide, ce sont des exemples d'effets secondaires Python et ne se produisent donc que pendant le traçage.

@tf.function
def buggy_consume_next(iterator):
  tf.print("Value:", next(iterator))

iterator = iter([1, 2, 3])
buggy_consume_next(iterator)
# This reuses the first value from the iterator, rather than consuming the next value.
buggy_consume_next(iterator)
buggy_consume_next(iterator)
Value: 1
Value: 1
Value: 1

Tout comme TensorFlow a un tf.TensorArray spécialisé pour les constructions de liste, il a un tf.data.Iterator spécialisé pour les constructions d'itération. Voir la section sur les transformations AutoGraph pour un aperçu. De plus, l'API tf.data peut aider à implémenter des modèles de générateur :

@tf.function
def good_consume_next(iterator):
  # This is ok, iterator is a tf.data.Iterator
  tf.print("Value:", next(iterator))

ds = tf.data.Dataset.from_tensor_slices([1, 2, 3])
iterator = iter(ds)
good_consume_next(iterator)
good_consume_next(iterator)
good_consume_next(iterator)
Value: 1
Value: 2
Value: 3

Suppression de tf.Variables entre Function appels de Function

Une autre erreur que vous pouvez rencontrer est une variable récupérée par la mémoire. ConcreteFunction s ne conservent que les WeakRefs des variables sur lesquelles elles se ferment, vous devez donc conserver une référence à toutes les variables.

external_var = tf.Variable(3)
@tf.function
def f(x):
  return x * external_var

traced_f = f.get_concrete_function(4)
print("Calling concrete function...")
print(traced_f(4))

# The original variable object gets garbage collected, since there are no more
# references to it.
external_var = tf.Variable(4)
print()
print("Calling concrete function after garbage collecting its closed Variable...")
with assert_raises(tf.errors.FailedPreconditionError):
  traced_f(4)
Calling concrete function...
tf.Tensor(12, shape=(), dtype=int32)

Calling concrete function after garbage collecting its closed Variable...
Caught expected exception 
  <class 'tensorflow.python.framework.errors_impl.FailedPreconditionError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-9a93d2e07632>", line 16, in <module>
    traced_f(4)
tensorflow.python.framework.errors_impl.FailedPreconditionError:  Could not find variable _AnonymousVar3. This could mean that the variable has been deleted. In TF1, it can also mean the variable is uninitialized. Debug info: container=localhost, status=Not found: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
     [[node ReadVariableOp (defined at <ipython-input-1-9a93d2e07632>:4) ]] [Op:__inference_f_782]

Function call stack:
f

Problèmes connus

Si votre Function n'est pas évaluée correctement, l'erreur peut s'expliquer par ces problèmes connus qui devraient être résolus à l'avenir.

En fonction des variables globales et libres de Python

Function crée une nouvelle ConcreteFunction lorsqu'elle est appelée avec une nouvelle valeur d'un argument Python. Cependant, il ne le fait pas pour la fermeture Python, les globals ou les nonlocals de cette Function . Si leur valeur change entre les appels à la Function , la Function utilisera toujours les valeurs qu'elle avait lorsqu'elle a été tracée. Ceci est différent du fonctionnement des fonctions Python classiques.

Pour cette raison, nous recommandons un style de programmation fonctionnel qui utilise des arguments au lieu de fermer sur des noms externes.

@tf.function
def buggy_add():
  return 1 + foo

@tf.function
def recommended_add(foo):
  return 1 + foo

foo = 1
print("Buggy:", buggy_add())
print("Correct:", recommended_add(foo))
Buggy: tf.Tensor(2, shape=(), dtype=int32)
Correct: tf.Tensor(2, shape=(), dtype=int32)
print("Updating the value of `foo` to 100!")
foo = 100
print("Buggy:", buggy_add())  # Did not change!
print("Correct:", recommended_add(foo))
Updating the value of `foo` to 100!
Buggy: tf.Tensor(2, shape=(), dtype=int32)
Correct: tf.Tensor(101, shape=(), dtype=int32)

Vous pouvez fermer les noms externes, tant que vous ne mettez pas à jour leurs valeurs.

Selon les objets Python

La recommandation de passer des objets Python en tant qu'arguments dans tf.function a un certain nombre de problèmes connus, qui devraient être corrigés à l'avenir. En général, vous pouvez compter sur un traçage cohérent si vous utilisez une primitive Python ou tf.nest structure compatible tf.nest comme argument ou si vous passez une instance différente d'un objet dans une Function . Cependant, Function ne créera pas de nouvelle trace lorsque vous passerez le même objet et ne modifiera que ses attributs .

class SimpleModel(tf.Module):
  def __init__(self):
    # These values are *not* tf.Variables.
    self.bias = 0.
    self.weight = 2.

@tf.function
def evaluate(model, x):
  return model.weight * x + model.bias

simple_model = SimpleModel()
x = tf.constant(10.)
print(evaluate(simple_model, x))
tf.Tensor(20.0, shape=(), dtype=float32)
print("Adding bias!")
simple_model.bias += 5.0
print(evaluate(simple_model, x))  # Didn't change :(
Adding bias!
tf.Tensor(20.0, shape=(), dtype=float32)

L'utilisation de la même Function pour évaluer l'instance mise à jour du modèle sera boguée car le modèle mis à jour a la même clé de cache que le modèle d'origine.

Pour cette raison, nous vous recommandons d'écrire votre Function pour éviter de dépendre des attributs d'objet mutables ou de créer de nouveaux objets.

Si cela n'est pas possible, une solution de contournement consiste à créer de nouvelles Function chaque fois que vous modifiez votre objet pour forcer le retraçage :

def evaluate(model, x):
  return model.weight * x + model.bias

new_model = SimpleModel()
evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x)
# Don't pass in `new_model`, `Function` already captured its state during tracing.
print(evaluate_no_bias(x))
tf.Tensor(20.0, shape=(), dtype=float32)
print("Adding bias!")
new_model.bias += 5.0
# Create new Function and ConcreteFunction since you modified new_model.
evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x)
print(evaluate_with_bias(x)) # Don't pass in `new_model`.
Adding bias!
tf.Tensor(25.0, shape=(), dtype=float32)

Comme le retraçage peut être coûteux , vous pouvez utiliser tf.Variable s comme attributs d'objet, qui peuvent être mutés (mais pas modifiés, attention !) pour un effet similaire sans avoir besoin de retracer.

class BetterModel:

  def __init__(self):
    self.bias = tf.Variable(0.)
    self.weight = tf.Variable(2.)

@tf.function
def evaluate(model, x):
  return model.weight * x + model.bias

better_model = BetterModel()
print(evaluate(better_model, x))
tf.Tensor(20.0, shape=(), dtype=float32)
print("Adding bias!")
better_model.bias.assign_add(5.0)  # Note: instead of better_model.bias += 5
print(evaluate(better_model, x))  # This works!
Adding bias!
tf.Tensor(25.0, shape=(), dtype=float32)

Création de tf.Variables

Function ne prend en charge la création de variables qu'une seule fois, lors du premier appel, puis leur réutilisation. Vous ne pouvez pas créer de tf.Variables dans de nouvelles traces. La création de nouvelles variables dans les appels suivants n'est actuellement pas autorisée, mais le sera à l'avenir.

Exemple:

@tf.function
def f(x):
  v = tf.Variable(1.0)
  return v

with assert_raises(ValueError):
  f(1.0)
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-8a0913e250e0>", line 7, in <module>
    f(1.0)
ValueError: in user code:

    <ipython-input-1-8a0913e250e0>:3 f  *
        v = tf.Variable(1.0)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:262 __call__  **
        return cls._variable_v2_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:256 _variable_v2_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:769 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

Vous pouvez créer des variables à l'intérieur d'une Function tant que ces variables ne sont créées que la première fois que la fonction est exécutée.

class Count(tf.Module):
  def __init__(self):
    self.count = None

  @tf.function
  def __call__(self):
    if self.count is None:
      self.count = tf.Variable(0)
    return self.count.assign_add(1)

c = Count()
print(c())
print(c())
tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Utilisation avec plusieurs optimiseurs Keras

Vous pouvez rencontrer ValueError: tf.function-decorated function tried to create variables on non-first call. lors de l'utilisation de plusieurs optimiseurs Keras avec une fonction tf.function . Cette erreur se produit car les optimiseurs créent en interne des tf.Variables lorsqu'ils appliquent des dégradés pour la première fois.

opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)
opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)

@tf.function
def train_step(w, x, y, optimizer):
   with tf.GradientTape() as tape:
       L = tf.reduce_sum(tf.square(w*x - y))
   gradients = tape.gradient(L, [w])
   optimizer.apply_gradients(zip(gradients, [w]))

w = tf.Variable(2.)
x = tf.constant([-1.])
y = tf.constant([2.])

train_step(w, x, y, opt1)
print("Calling `train_step` with different optimizer...")
with assert_raises(ValueError):
  train_step(w, x, y, opt2)
Calling `train_step` with different optimizer...
Caught expected exception 
  <class 'ValueError'>:
Traceback (most recent call last):
  File "<ipython-input-1-73d0ca52e838>", line 8, in assert_raises
    yield
  File "<ipython-input-1-d3d3937dbf1a>", line 18, in <module>
    train_step(w, x, y, opt2)
ValueError: in user code:

    <ipython-input-1-d3d3937dbf1a>:9 train_step  *
        optimizer.apply_gradients(zip(gradients, [w]))
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:636 apply_gradients  **
        self._create_all_weights(var_list)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:821 _create_all_weights
        _ = self.iterations
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:828 __getattribute__
        return super(OptimizerV2, self).__getattribute__(name)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:988 iterations
        aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:1194 add_weight
        aggregation=aggregation)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:815 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer_utils.py:139 make_variable
        shape=variable_shape if variable_shape else None)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:260 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:769 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

Si vous devez modifier l'optimiseur pendant la formation, une solution de contournement consiste à créer une nouvelle Function pour chaque optimiseur, en appelant directement la ConcreteFunction .

opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)
opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)

# Not a tf.function.
def train_step(w, x, y, optimizer):
   with tf.GradientTape() as tape:
       L = tf.reduce_sum(tf.square(w*x - y))
   gradients = tape.gradient(L, [w])
   optimizer.apply_gradients(zip(gradients, [w]))

w = tf.Variable(2.)
x = tf.constant([-1.])
y = tf.constant([2.])

# Make a new Function and ConcreteFunction for each optimizer.
train_step_1 = tf.function(train_step).get_concrete_function(w, x, y, opt1)
train_step_2 = tf.function(train_step).get_concrete_function(w, x, y, opt2)
for i in range(10):
  if i % 2 == 0:
    train_step_1(w, x, y) # `opt1` is not used as a parameter. 
  else:
    train_step_2(w, x, y) # `opt2` is not used as a parameter.

Utilisation avec plusieurs modèles Keras

Vous pouvez également rencontrer ValueError: tf.function-decorated function tried to create variables on non-first call. lors du passage de différentes instances de modèle à la même Function .

Cette erreur se produit car les modèles Keras (dont la forme d'entrée n'est pas définie ) et les couches Keras créent des tf.Variables s lors de leur premier appel. Vous essayez peut-être d'initialiser ces variables dans une Function , qui a déjà été appelée. Pour éviter cette erreur, essayez d'appeler model.build(input_shape) pour initialiser tous les poids avant d'entraîner le modèle.

Lectures complémentaires

Pour savoir comment exporter et charger une Function , consultez le guide SavedModel . Pour en savoir plus sur les optimisations de graphes effectuées après le traçage, consultez le guide Grappler . Pour savoir comment optimiser votre pipeline de données et profiler votre modèle, consultez le guide Profiler .