Den @tf.function Decorator von TensorFlow verstehen

Einleitung

Durch die Verbesserung der Leistung einer Trainingsschleife können beim Trainieren von Modellen für maschinelles Lernen Stunden an Rechenzeit eingespart werden. Eine Möglichkeit, die Leistung von TensorFlow-Code zu verbessern, ist die Verwendung von tf.function() decorator – eine einfache, einzeilige Änderung, die Ihre Funktionen erheblich schneller ausführen kann.

In dieser kurzen Anleitung erklären wir, wie tf.function() verbessert die Leistung und werfen Sie einen Blick auf einige Best Practices.

Python-Dekorateure und tf.Funktion()

In Python ist ein Decorator eine Funktion, die das Verhalten anderer Funktionen modifiziert. Angenommen, Sie rufen die folgende Funktion in einer Notizbuchzelle auf:

import tensorflow as tf

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

def some_costly_computation(x):
    aux = tf.eye(100, dtype=tf.dtypes.float32)
    result = tf.zeros(100, dtype = tf.dtypes.float32)
    for i in range(1,100):
        aux = tf.matmul(x,aux)/i
        result = result + aux
    return result

%timeit some_costly_computation(x)

16.2 ms ± 103 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Wenn wir jedoch die kostspielige Funktion in a übergeben tf.function():

quicker_computation = tf.function(some_costly_computation)
%timeit quicker_computation(x)

Wir bekommen quicker_computation() – eine neue Funktion, die viel schneller arbeitet als die vorherige:

4.99 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Damit tf.function() modifiziert some_costly_computation() und gibt die aus quicker_computation() Funktion. Dekorateure modifizieren auch Funktionen, also war es natürlich zu machen tf.function() auch ein Dekorateur.

Die Verwendung der Decorator-Notation ist dasselbe wie das Aufrufen tf.function(function):

@tf.function
def quick_computation(x):
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

%timeit quick_computation(x)

5.09 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Wie funktioniert tf.function() Arbeit?

Wie kommt es, dass wir bestimmte Funktionen 2-3x schneller ausführen können?

TensorFlow-Code kann in zwei Modi ausgeführt werden: Eifermodus und Grafikmodus. Der Eager-Modus ist die interaktive Standardmethode zum Ausführen von Code: Jedes Mal, wenn Sie eine Funktion aufrufen, wird sie ausgeführt.

Der Grafikmodus ist jedoch etwas anders. Im Diagrammmodus erstellt TensorFlow vor dem Ausführen der Funktion einen Berechnungsgraphen, bei dem es sich um eine Datenstruktur handelt, die die zum Ausführen der Funktion erforderlichen Operationen enthält. Das Berechnungsdiagramm ermöglicht es TensorFlow, die Berechnungen zu vereinfachen und Möglichkeiten zur Parallelisierung zu finden. Das Diagramm isoliert die Funktion auch vom darüber liegenden Python-Code, sodass sie effizient auf vielen verschiedenen Geräten ausgeführt werden kann.

Eine Funktion dekoriert mit @tf.function wird in zwei Schritten ausgeführt:

1. Im ersten Schritt führt TensorFlow den Python-Code für die Funktion aus und kompiliert ein Berechnungsdiagramm, wodurch die Ausführung aller TensorFlow-Operationen verzögert wird.
2. Anschließend wird der Berechnungsgraph ausgeführt.

Der erste Schritt wird übersprungen, wenn kein neuer Berechnungsgraph erstellt werden muss. Dies verbessert die Leistung der Funktion, bedeutet aber auch, dass die Funktion nicht wie normaler Python-Code ausgeführt wird (in dem jede ausführbare Zeile ausgeführt wird). Ändern wir zum Beispiel unsere vorherige Funktion:

@tf.function
def quick_computation(x):
  print('Only prints the first time!')
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation(x)
quick_computation(x)

Das führt zu:

Only prints the first time!

Das print() wird nur einmal während des Ablaufverfolgungsschritts ausgeführt, wenn regulärer Python-Code ausgeführt wird. Die nächsten Aufrufe der Funktion führen nur TenforFlow-Operationen aus dem Berechnungsgraphen (TensorFlow-Operationen) aus.

Wenn wir jedoch verwenden tf.print() stattdessen:

@tf.function
def quick_computation_with_print(x):
  tf.print("Prints every time!")
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation_with_print(x)
quick_computation_with_print(x)

Prints every time!
Prints every time!

TensorFlow beinhaltet tf.print() in seinem Berechnungsdiagramm, da es sich um eine TensorFlow-Operation handelt – nicht um eine reguläre Python-Funktion.

Best Practices mit @tf.function

Schreiben von Code mit TensorFlow-Operationen

Wie wir gerade gezeigt haben, werden einige Teile des Codes vom Berechnungsgraphen ignoriert. Dies macht es schwierig, das Verhalten der Funktion vorherzusagen, wenn mit „normalem“ Python-Code codiert wird, wie wir gerade gesehen haben print(). Es ist besser, Ihre Funktion gegebenenfalls mit TensorFlow-Vorgängen zu codieren, um unerwartetes Verhalten zu vermeiden.

Zum Beispiel for und while Schleifen können in die äquivalente TensorFlow-Schleife konvertiert werden oder nicht. Daher ist es besser, wenn möglich, Ihre „for“-Schleife als vektorisierte Operation zu schreiben. Dadurch wird die Leistung Ihres Codes verbessert und sichergestellt, dass Ihre Funktion korrekt nachverfolgt wird.

Betrachten Sie als Beispiel Folgendes:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

@tf.function
def function_with_for(x):
    summ = float(0)
    for row in x:
      summ = summ + tf.reduce_mean(row)
    return summ

@tf.function
def vectorized_function(x):
  result = tf.reduce_mean(x, axis=0)
  return tf.reduce_sum(result)


print(function_with_for(x))
print(vectorized_function(x))

%timeit function_with_for(x)
%timeit vectorized_function(x)

tf.Tensor(0.672811, shape=(), dtype=float32)
tf.Tensor(0.67281103, shape=(), dtype=float32)
1.58 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
440 µs ± 8.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Der Code mit den TensorFlow-Operationen ist erheblich schneller.

Vermeiden Sie Verweise auf globale Variablen

Betrachten Sie den folgenden Code:

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = 2

@tf.function
def power(x):
  return tf.pow(x,y)

print(power(x))

y = 3

print(power(x))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(4.0, shape=(), dtype=float32)

Das erste Mal die dekorierte Funktion power() aufgerufen wurde, war der Ausgabewert der erwartete Wert 4. Beim zweiten Mal ignorierte die Funktion jedoch, dass der Wert von y wurde verändert. Dies geschieht, weil der Wert der globalen Python-Variablen nach der Ablaufverfolgung für die Funktion eingefroren wird.

Ein besserer Weg wäre zu verwenden tf.Variable() für alle Ihre Variablen und übergeben Sie beide als Argumente an Ihre Funktion.

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = tf.Variable(2, dtype = tf.dtypes.float32)

@tf.function
def power(x,y):
  return tf.pow(x,y)

print(power(x,y))

y.assign(3)

print(power(x,y))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(8.0, shape=(), dtype=float32)

Fehlerbeseitigung [E-Mail geschützt] _s

Im Allgemeinen möchten Sie Ihre Funktion im Eifer-Modus debuggen und sie dann mit dekorieren @tf.function nachdem Ihr Code korrekt ausgeführt wurde, da die Fehlermeldungen im Eifer-Modus informativer sind.

Einige häufige Probleme sind Tippfehler und Formfehler. Typfehler treten auf, wenn der Typ der an einer Operation beteiligten Variablen nicht übereinstimmt:

x = tf.Variable(1, dtype = tf.dtypes.float32)
y = tf.Variable(1, dtype = tf.dtypes.int32)

z = tf.add(x,y)

InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a float tensor but is a int32 tensor [Op:AddV2]

Typfehler schleichen sich leicht ein und können leicht behoben werden, indem eine Variable in einen anderen Typ umgewandelt wird:

y = tf.cast(y, tf.dtypes.float32)
z = tf.add(x, y) 
tf.print(z) 

Formfehler treten auf, wenn Ihre Tensoren nicht die Form haben, die Ihre Operation erfordert:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)
y = tf.random.uniform(shape=[1, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

z = tf.matmul(x,y)

InvalidArgumentError: Matrix size-incompatible: In[0]: [100,100], In[1]: [1,100] [Op:MatMul]

Ein praktisches Tool zum Beheben beider Arten von Fehlern ist der interaktive Python-Debugger, mit dem Sie automatisch in einem Jupyter-Notebook aufrufen können %pdb. Damit können Sie Ihre Funktion codieren und durch einige gängige Anwendungsfälle ausführen. Wenn ein Fehler auftritt, wird eine interaktive Eingabeaufforderung geöffnet. Mit dieser Eingabeaufforderung können Sie die Abstraktionsebenen in Ihrem Code auf und ab gehen und die Werte, Typen und Formen Ihrer TensorFlow-Variablen überprüfen.

Zusammenfassung

Wir haben gesehen, wie TensorFlow's tf.function() Ihre Funktion effizienter macht und wie die @tf.function decorator wendet die Funktion auf Ihre eigene an.

Diese Beschleunigung ist nützlich bei Funktionen, die viele Male aufgerufen werden, wie z. B. benutzerdefinierte Trainingsschritte für Modelle für maschinelles Lernen.