Forståelse af TensorFlows @tf.function Decorator

Introduktion

Forbedring af ydeevnen af ​​en træningsløkke kan spare timers computertid, når du træner maskinlæringsmodeller. En af måderne til at forbedre ydeevnen af ​​TensorFlow-kode er at bruge tf.function() decorator – en enkel ændring på én linje, der kan få dine funktioner til at køre markant hurtigere.

I denne korte guide vil vi forklare hvordan tf.function() forbedrer ydeevnen og tag et kig på nogle bedste fremgangsmåder.

Python dekoratører og tf.function()

I Python er en dekoratør en funktion, der ændrer adfærden af ​​andre funktioner. Antag for eksempel, at du kalder følgende funktion i en notesbogscelle:

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)

Men hvis vi overfører den kostbare funktion til en tf.function():

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

Vi får quicker_computation() – en ny funktion, der udfører meget hurtigere end den forrige:

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

Så, tf.function() modificerer some_costly_computation() og udsender quicker_computation() fungere. Dekoratører ændrer også funktioner, så det var naturligt at lave tf.function() også en dekoratør.

At bruge dekorationsnotationen er det samme som at kalde 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)

Hvordan virker det? tf.function() Arbejde?

Hvorfor kan vi få visse funktioner til at køre 2-3 gange hurtigere?

TensorFlow-kode kan køres i to tilstande: ivrig tilstand , graftilstand. Ivrig tilstand er den interaktive standard måde at køre kode på: hver gang du kalder en funktion, udføres den.

Graph mode er dog en smule anderledes. I graftilstand, før funktionen udføres, opretter TensorFlow en beregningsgraf, som er en datastruktur, der indeholder de operationer, der kræves for at udføre funktionen. Beregningsgrafen giver TensorFlow mulighed for at forenkle beregningerne og finde muligheder for parallelisering. Grafen isolerer også funktionen fra den overliggende Python-kode, så den kan køres effektivt på mange forskellige enheder.

En funktion dekoreret med @tf.function udføres i to trin:

1. I det første trin udfører TensorFlow Python-koden for funktionen og kompilerer en beregningsgraf, hvilket forsinker udførelsen af ​​enhver TensorFlow-operation.
2. Bagefter køres beregningsgrafen.

Det første trin springes over, hvis der ikke er behov for at oprette en ny beregningsgraf. Dette forbedrer funktionens ydeevne, men betyder også, at funktionen ikke vil køre som almindelig Python-kode (hvori hver eksekverbar linje udføres). Lad os for eksempel ændre vores tidligere 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)

Dette resulterer i:

Only prints the first time!

print() udføres kun én gang under sporingstrinnet, hvilket er når almindelig Python-kode køres. De næste kald til funktionen udfører kun TenforFlow-operationer fra beregningsgrafen (TensorFlow-operationer).

Men hvis vi bruger tf.print() i stedet:

@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 inkluderer tf.print() i sin beregningsgraf, da det er en TensorFlow-operation – ikke en almindelig Python-funktion.

Bedste praksis med @tf.function

Skrive kode med TensorFlow Operations

Som vi lige har vist, ignoreres nogle dele af koden af ​​beregningsgrafen. Dette gør det svært at forudsige funktionens adfærd ved kodning med "normal" Python-kode, som vi lige har set med print(). Det er bedre at kode din funktion med TensorFlow-operationer, når det er relevant for at undgå uventet adfærd.

For eksempel, for , while sløjfer kan eller kan ikke konverteres til den tilsvarende TensorFlow-løkke. Derfor er det bedre at skrive din "for"-løkke som en vektoriseret operation, hvis det er muligt. Dette vil forbedre ydeevnen af ​​din kode og sikre, at din funktion spores korrekt.

Som et eksempel kan du overveje følgende:

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)

Koden med TensorFlow-operationerne er betydeligt hurtigere.

Undgå referencer til globale variabler

Overvej følgende kode:

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)

Første gang den dekorerede funktion power() blev kaldt, var outputværdien den forventede 4. Men anden gang ignorerede funktionen, at værdien af y blev ændret. Dette sker, fordi værdien af ​​Python globale variabler er frosset for funktionen efter sporing.

En bedre måde ville være at bruge tf.Variable() for alle dine variabler og videregive begge som argumenter til din 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)

Debugging [e-mail beskyttet]_s

Generelt vil du gerne fejlsøge din funktion i ivrig tilstand og derefter dekorere dem med @tf.function efter din kode kører korrekt, fordi fejlmeddelelserne i ivrig tilstand er mere informative.

Nogle almindelige problemer er typefejl og formfejl. Typefejl opstår, når der er et misforhold i typen af ​​de variable involveret i en operation:

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]

Typefejl sniger sig nemt ind og kan nemt rettes ved at caste en variabel til en anden type:

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

Formfejl opstår, når dine tensorer ikke har den form, din operation kræver:

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]

Et praktisk værktøj til at rette begge slags fejl er den interaktive Python debugger, som du kan kalde automatisk i en Jupyter Notebook vha. %pdb. Ved at bruge det kan du kode din funktion og køre den gennem nogle almindelige use cases. Hvis der er en fejl, åbnes en interaktiv prompt. Denne prompt giver dig mulighed for at gå op og ned i abstraktionslagene i din kode og kontrollere værdierne, typerne og formerne af dine TensorFlow-variabler.

Konklusion

Vi har set, hvordan TensorFlow er tf.function() gør din funktion mere effektiv, og hvordan @tf.function dekoratør anvender funktionen til din egen.

Denne fremskyndelse er nyttig i funktioner, der vil blive kaldt mange gange, såsom tilpassede træningstrin til maskinlæringsmodeller.