Forstå TensorFlows @tf.function-dekorator

Introduksjon

Å forbedre ytelsen til en treningssløyfe kan spare timer med databehandlingstid når du trener maskinlæringsmodeller. En av måtene å forbedre ytelsen til TensorFlow-koden på er å bruke tf.function() decorator – en enkel endring på én linje som kan få funksjonene dine til å kjøre betydelig raskere.

I denne korte veiledningen vil vi forklare hvordan tf.function() forbedrer ytelsen og ta en titt på noen beste fremgangsmåter.

Python-dekoratører og tf.function()

I Python er en dekoratør en funksjon som endrer oppførselen til andre funksjoner. Anta for eksempel at du kaller følgende funksjon i en bærbar celle:

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)

Imidlertid, hvis vi overfører den kostbare funksjonen til en tf.function():

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

Vi får quicker_computation() – en ny funksjon som utfører mye raskere enn den forrige:

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

Så, tf.function() modifiserer some_costly_computation() og gir ut quicker_computation() funksjon. Dekoratører endrer også funksjoner, så det var naturlig å lage tf.function() en dekoratør også.

Å bruke dekorasjonsnotasjonen er det samme som å ringe 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 gjør tf.function() Arbeid?

Hvordan kan vi få visse funksjoner til å kjøre 2-3 ganger raskere?

TensorFlow-kode kan kjøres i to moduser: ivrig modus og grafmodus. Ivrig-modus er standard, interaktiv måte å kjøre kode på: hver gang du kaller en funksjon, blir den utført.

Grafmodus er imidlertid litt annerledes. I grafmodus, før du utfører funksjonen, lager TensorFlow en beregningsgraf, som er en datastruktur som inneholder operasjonene som kreves for å utføre funksjonen. Beregningsgrafen lar TensorFlow forenkle beregningene og finne muligheter for parallellisering. Grafen isolerer også funksjonen fra den overliggende Python-koden, slik at den kan kjøres effektivt på mange forskjellige enheter.

En funksjon dekorert med @tf.function utføres i to trinn:

1. I det første trinnet utfører TensorFlow Python-koden for funksjonen og kompilerer en beregningsgraf, noe som forsinker utførelsen av enhver TensorFlow-operasjon.
2. Etterpå kjøres beregningsgrafen.

Det første trinnet hoppes over hvis det ikke er behov for å lage en ny beregningsgraf. Dette forbedrer ytelsen til funksjonen, men betyr også at funksjonen ikke vil kjøre som vanlig Python-kode (der hver kjørbar linje kjøres). La oss for eksempel endre vår forrige funksjon:

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

De print() utføres kun én gang under sporingstrinnet, som er når vanlig Python-kode kjøres. De neste kallene til funksjonen utfører kun TenforFlow-operasjoner fra beregningsgrafen (TensorFlow-operasjoner).

Imidlertid, hvis vi bruker 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 beregningsgrafen siden det er en TensorFlow-operasjon – ikke en vanlig Python-funksjon.

Beste praksis med @tf.function

Skrive kode med TensorFlow-operasjoner

Som vi nettopp har vist, ignoreres noen deler av koden av beregningsgrafen. Dette gjør det vanskelig å forutsi funksjonen til funksjonen ved koding med "normal" Python-kode, som vi nettopp har sett med print(). Det er bedre å kode funksjonen din med TensorFlow-operasjoner når det er aktuelt for å unngå uventet oppførsel.

Eksempelvis for og while løkker kan eller kan ikke konverteres til tilsvarende TensorFlow-løkke. Derfor er det bedre å skrive "for"-løkken din som en vektorisert operasjon, hvis mulig. Dette vil forbedre ytelsen til koden din og sikre at funksjonen spores riktig.

Som et eksempel kan du vurdere 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-operasjonene er betydelig raskere.

Unngå referanser til globale variabler

Vurder 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 dekorerte funksjonen power() ble kalt, var utgangsverdien den forventede 4. Men den andre gangen ignorerte funksjonen at verdien av y ble endret. Dette skjer fordi verdien av Python globale variabler fryses for funksjonen etter sporing.

En bedre måte ville være å bruke tf.Variable() for alle variablene dine og send begge som argumenter til funksjonen din.

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-postbeskyttet]_s

Generelt vil du feilsøke funksjonen din i ivrig modus, og deretter dekorere dem med @tf.function etter at koden kjører riktig fordi feilmeldingene i ivrig modus er mer informative.

Noen vanlige problemer er typefeil og formfeil. Typefeil oppstår når det er uoverensstemmelse i typen av variablene som er involvert i en operasjon:

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]

Typefeil kommer lett inn, og kan enkelt fikses ved å caste en variabel til en annen type:

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

Formfeil oppstår når tensorene dine ikke har formen operasjonen krever:

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 verktøy for å fikse begge typer feil er den interaktive Python-feilsøkeren, som du kan kalle opp automatisk i en Jupyter Notebook ved å bruke %pdb. Ved å bruke det kan du kode funksjonen din og kjøre den gjennom noen vanlige brukstilfeller. Hvis det er en feil, åpnes en interaktiv melding. Denne ledeteksten lar deg gå opp og ned abstraksjonslagene i koden din og sjekke verdiene, typene og formene til TensorFlow-variablene dine.

konklusjonen

Vi har sett hvordan TensorFlow er tf.function() gjør funksjonen din mer effektiv, og hvordan @tf.function dekoratør bruker funksjonen til din egen.

Denne hastigheten er nyttig i funksjoner som vil bli kalt mange ganger, for eksempel tilpassede opplæringstrinn for maskinlæringsmodeller.