TensorFlow's @tf.function Decorator begrijpen

Introductie

Het verbeteren van de prestaties van een trainingslus kan uren aan rekentijd besparen bij het trainen van machine learning-modellen. Een van de manieren om de prestaties van TensorFlow-code te verbeteren is het gebruik van de tf.function() decorateur – een eenvoudige wijziging in één regel waarmee uw functies aanzienlijk sneller kunnen werken.

In deze korte handleiding leggen we uit hoe tf.function() verbetert de prestaties en bekijk enkele best practices.

Python-decorateurs en tf.function()

In Python is een decorateur een functie die het gedrag van andere functies wijzigt. Stel dat u bijvoorbeeld de volgende functie aanroept in een notebookcel:

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)

Als we echter de kostbare functie doorgeven aan a tf.function():

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

We krijgen quicker_computation() – een nieuwe functie die veel sneller presteert dan de vorige:

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

Dus, tf.function() wijzigt some_costly_computation() en voert de quicker_computation() functie. Decorateurs passen ook functies aan, dus het was vanzelfsprekend om te maken tf.function() ook een decorateur.

Het gebruik van de decorateurnotatie is hetzelfde als bellen 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)

Hoe werkt tf.function() Werk?

Hoe komt het dat we bepaalde functies 2-3x sneller kunnen laten werken?

TensorFlow-code kan in twee modi worden uitgevoerd: enthousiaste modus en grafiek modus. Eager-modus is de standaard, interactieve manier om code uit te voeren: elke keer dat u een functie aanroept, wordt deze uitgevoerd.

De grafiekmodus is echter een beetje anders. In de grafiekmodus maakt TensorFlow, voordat de functie wordt uitgevoerd, een berekeningsgrafiek, een datastructuur die de bewerkingen bevat die nodig zijn voor het uitvoeren van de functie. Met de berekeningsgrafiek kan TensorFlow de berekeningen vereenvoudigen en mogelijkheden voor parallellisatie vinden. De grafiek isoleert de functie ook van de bovenliggende Python-code, waardoor deze efficiënt op veel verschillende apparaten kan worden uitgevoerd.

Een functie versierd met @tf.function wordt in twee stappen uitgevoerd:

1. In de eerste stap voert TensorFlow de Python-code voor de functie uit en stelt een berekeningsgrafiek samen, waardoor de uitvoering van een TensorFlow-bewerking wordt vertraagd.
2. Daarna wordt de berekeningsgrafiek uitgevoerd.

De eerste stap wordt overgeslagen als het niet nodig is een nieuwe berekeningsgrafiek te maken. Dit verbetert de prestaties van de functie, maar betekent ook dat de functie niet zal worden uitgevoerd zoals gewone Python-code (waarin elke uitvoerbare regel wordt uitgevoerd). Laten we bijvoorbeeld onze vorige functie aanpassen:

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

Dit resulteert in:

Only prints the first time!

De print() wordt slechts één keer uitgevoerd tijdens de traceringsstap, namelijk wanneer reguliere Python-code wordt uitgevoerd. De volgende aanroepen van de functie voeren alleen TenforFlow-bewerkingen uit de berekeningsgrafiek uit (TensorFlow-bewerkingen).

Echter, als we gebruiken tf.print() in plaats daarvan:

@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 omvat tf.print() in de berekeningsgrafiek omdat het een TensorFlow-bewerking is en geen gewone Python-functie.

Beste praktijken met @tf.function

Code schrijven met TensorFlow Operations

Zoals we zojuist hebben laten zien, worden sommige delen van de code genegeerd door de berekeningsgrafiek. Dit maakt het moeilijk om het gedrag van de functie te voorspellen bij het coderen met “normale” Python-code, zoals we zojuist hebben gezien met print(). Het is beter om uw functie te coderen met TensorFlow-bewerkingen, indien van toepassing, om onverwacht gedrag te voorkomen.

Bijvoorbeeld, for en while lussen kunnen al dan niet worden omgezet in de equivalente TensorFlow-lus. Daarom is het beter om uw “for”-lus, indien mogelijk, als een gevectoriseerde bewerking te schrijven. Dit verbetert de prestaties van uw code en zorgt ervoor dat uw functie correct wordt getraceerd.

Beschouw als voorbeeld het volgende:

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)

De code met de TensorFlow-bewerkingen is aanzienlijk sneller.

Vermijd verwijzingen naar globale variabelen

Beschouw de volgende 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)

De eerste keer de versierde functie power() werd aangeroepen, was de uitvoerwaarde de verwachte 4. De tweede keer negeerde de functie echter dat de waarde van y was veranderd. Dit gebeurt omdat de waarde van de globale Python-variabelen na tracering voor de functie wordt bevroren.

Een betere manier zou zijn om te gebruiken tf.Variable() voor al uw variabelen en geef beide als argumenten door aan uw functie.

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)

Debugging [e-mail beveiligd]_s

Over het algemeen wil je je functie in de gretige modus debuggen en ze vervolgens versieren met @tf.function nadat uw code correct is uitgevoerd, omdat de foutmeldingen in de enthousiaste modus informatiever zijn.

Enkele veel voorkomende problemen zijn typefouten en vormfouten. Typefouten treden op als er een mismatch is in het type van de variabelen die bij een bewerking betrokken zijn:

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]

Typefouten sluipen er gemakkelijk in en kunnen eenvoudig worden verholpen door een variabele naar een ander type te casten:

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

Vormfouten treden op wanneer uw tensoren niet de vorm hebben die uw operatie vereist:

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]

Een handig hulpmiddel voor het oplossen van beide soorten fouten is de interactieve Python-debugger, die u automatisch in een Jupyter Notebook kunt oproepen met behulp van %pdb. Hiermee kunt u uw functie coderen en enkele veelvoorkomende gebruiksscenario's doorlopen. Als er een fout optreedt, wordt een interactieve prompt geopend. Met deze prompt kunt u de abstractielagen in uw code op en neer gaan en de waarden, typen en vormen van uw TensorFlow-variabelen controleren.

Conclusie

We hebben gezien hoe TensorFlow werkt tf.function() uw functie efficiënter maakt, en hoe de @tf.function decorateur past de functie toe op uw eigen functie.

Deze versnelling is handig bij functies die vaak zullen worden aangeroepen, zoals aangepaste trainingsstappen voor machine learning-modellen.