A TensorFlow @tf.function Decorator megértése

Bevezetés

Egy betanítási hurok teljesítményének javítása órákon át tartó számítási időt takaríthat meg a gépi tanulási modellek betanítása során. A TensorFlow kód teljesítményének javításának egyik módja a tf.function() dekorátor – egy egyszerű, egysoros változtatás, amellyel lényegesen gyorsabban futhatnak a funkciók.

Ebben a rövid útmutatóban elmagyarázzuk, hogyan tf.function() javítja a teljesítményt, és vessen egy pillantást néhány bevált gyakorlatra.

Python Decorators és tf.function()

A Pythonban a dekorátor olyan függvény, amely más függvények viselkedését módosítja. Tegyük fel például, hogy meghívja a következő függvényt egy notebook cellában:

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)

Ha azonban a költséges függvényt átadjuk a tf.function():

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

Kapunk quicker_computation() – egy új funkció, amely sokkal gyorsabban működik, mint az előző:

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

Szóval, tf.function() módosítja some_costly_computation() és kiadja a quicker_computation() funkció. A dekoratőrök a funkciókat is módosítják, így természetes volt az elkészítése tf.function() lakberendező is.

A díszítő jelölés használata megegyezik a hívással 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)

Hogyan működik tf.function() Munka?

Miért tudunk bizonyos funkciókat 2-3-szor gyorsabbá tenni?

A TensorFlow kód két módban futtatható: mohó mód és a grafikon mód. Az Eager mód a kód futtatásának szabványos, interaktív módja: minden alkalommal, amikor meghív egy függvényt, az végrehajtódik.

A grafikon mód azonban egy kicsit más. Grafikon módban a függvény végrehajtása előtt a TensorFlow egy számítási gráfot készít, amely a függvény végrehajtásához szükséges műveleteket tartalmazó adatstruktúra. A számítási gráf lehetővé teszi a TensorFlow számára, hogy leegyszerűsítse a számításokat és megtalálja a párhuzamosítási lehetőségeket. A grafikon a függvényt a felül lévő Python-kódtól is elkülöníti, lehetővé téve annak hatékony futtatását számos különböző eszközön.

-vel díszített funkció @tf.function két lépésben hajtják végre:

1. Az első lépésben a TensorFlow végrehajtja a függvény Python-kódját, és összeállít egy számítási gráfot, késleltetve a TensorFlow műveletek végrehajtását.
2. Ezt követően lefut a számítási gráf.

Az első lépést a rendszer kihagyja, ha nincs szükség új számítási grafikon létrehozására. Ez javítja a függvény teljesítményét, de azt is jelenti, hogy a függvény nem fog úgy futni, mint a szokásos Python-kód (amelyben minden végrehajtható sor lefut). Például módosítsuk az előző függvényünket:

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

Ennek eredményeként:

Only prints the first time!

A print() csak egyszer fut le a nyomkövetési lépés során, vagyis amikor a szokásos Python-kód fut. A függvény következő hívásai csak TenforFlow műveleteket hajtanak végre a számítási gráfból (TensorFlow műveletek).

Ha azonban használjuk tf.print() helyette:

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

A TensorFlow tartalmazza tf.print() számítási grafikonján, mivel ez egy TensorFlow művelet – nem egy szokásos Python-függvény.

A legjobb gyakorlatok @tf.function

Kód írása TensorFlow műveletekkel

Amint az imént bemutattuk, a kód egyes részeit figyelmen kívül hagyja a számítási gráf. Ez megnehezíti a függvény viselkedésének megjósolását „normál” Python kóddal történő kódoláskor, ahogy az imént láthattuk print(). A váratlan viselkedés elkerülése érdekében jobb, ha a függvényt TensorFlow-műveletekkel kódolja, ha lehetséges.

Például, for és a while hurkok konvertálhatók egyenértékű TensorFlow hurokká, de nem. Ezért jobb, ha a „for” ciklust vektorizált műveletként írjuk, ha lehetséges. Ez javítja a kód teljesítményét, és biztosítja, hogy a függvény megfelelően nyomon követhető legyen.

Példaként vegye figyelembe a következőket:

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)

A TensorFlow műveletek kódja lényegesen gyorsabb.

Kerülje a globális változókra való hivatkozásokat

Vegye figyelembe a következő kódot:

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)

Az első alkalommal a díszített funkció power() meghívásra került, a kimeneti érték a várt 4 volt. Másodszor azonban a függvény figyelmen kívül hagyta, hogy az értéke a y megváltozott. Ez azért történik, mert a Python globális változók értéke a függvény számára lefagy a nyomkövetés után.

Jobb módszer lenne a használata tf.Variable() az összes változóhoz, és mindkettőt argumentumként adja át a függvénynek.

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)

Hibakeresés [e-mail védett]_s

Általában buzgó módban szeretné hibakeresni a függvényt, majd díszíteni @tf.function miután a kód megfelelően fut, mert a buzgó módban megjelenő hibaüzenetek informatívabbak.

Néhány gyakori probléma a típus- és alakhibák. Típushibák akkor fordulnak elő, ha a műveletben részt vevő változók típusa nem egyezik:

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]

A típushibák könnyen bekúsznak, és könnyen kijavíthatók egy változó más típusba öntésével:

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

Alakhibák akkor fordulnak elő, ha a tenzorok nem olyan alakúak, mint amilyet a művelet igényel:

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]

Mindkét típusú hiba kijavításának egyik kényelmes eszköze az interaktív Python hibakereső, amely automatikusan meghívható egy Jupyter Notebookban a %pdb. Ennek segítségével kódolhatja a függvényt, és futtathatja néhány gyakori használati eseten. Ha hiba történik, interaktív prompt nyílik meg. Ez a prompt lehetővé teszi a kód absztrakciós rétegeinek fel- és lefutását, valamint a TensorFlow-változók értékeinek, típusainak és alakjainak ellenőrzését.

Következtetés

Láttuk, hogyan működik a TensorFlow tf.function() hatékonyabbá teszi működését, és hogyan a @tf.function lakberendező alkalmazza a funkciót a sajátjára.

Ez a gyorsítás hasznos a sokszor meghívott függvényeknél, például a gépi tanulási modellek egyéni betanítási lépéseiben.