Razumevanje @tf.function Decoratorja TensorFlow

Predstavitev

Izboljšanje zmogljivosti vadbene zanke lahko prihrani ure računalniškega časa pri usposabljanju modelov strojnega učenja. Eden od načinov za izboljšanje učinkovitosti kode TensorFlow je uporaba tf.function() dekorater – preprosta sprememba v eni vrstici, s katero lahko vaše funkcije delujejo bistveno hitreje.

V tem kratkem vodniku vam bomo razložili, kako tf.function() izboljša delovanje in si oglejte nekaj najboljših praks.

Python dekoraterji in tf.function()

V Pythonu je dekorater funkcija, ki spreminja vedenje drugih funkcij. Recimo, da v celici beležnice pokličete naslednjo funkcijo:

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)

Vendar, če prenesemo drago funkcijo v a tf.function():

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

Dobimo quicker_computation() – nova funkcija, ki deluje veliko hitreje od prejšnje:

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

Torej, tf.function() spremenjene some_costly_computation() in izpiše quicker_computation() funkcijo. Dekoraterji prav tako spreminjajo funkcije, zato je bilo naravno narediti tf.function() tudi dekorater.

Uporaba notacije dekoraterja je enaka klicanju 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)

Kako tf.function() Delo?

Kako to, da lahko nekatere funkcije delujejo 2-3x hitreje?

Kodo TensorFlow je mogoče izvajati v dveh načinih: nestrpen način in grafični način. Način Eager je standardni, interaktivni način za zagon kode: vsakič, ko pokličete funkcijo, se ta izvede.

Način grafa pa je nekoliko drugačen. V grafičnem načinu TensorFlow pred izvedbo funkcije ustvari računski graf, ki je podatkovna struktura, ki vsebuje operacije, potrebne za izvedbo funkcije. Računski graf omogoča TensorFlowu, da poenostavi izračune in poišče priložnosti za paralelizacijo. Graf tudi izolira funkcijo od prekrivne kode Python, kar omogoča njeno učinkovito izvajanje na številnih različnih napravah.

Funkcija, okrašena z @tf.function se izvaja v dveh korakih:

1. V prvem koraku TensorFlow izvede kodo Python za funkcijo in sestavi računski graf, s čimer odloži izvajanje katere koli operacije TensorFlow.
2. Nato se zažene računski graf.

Prvi korak bo preskočen, če ni treba ustvariti novega računskega grafa. To izboljša delovanje funkcije, hkrati pa pomeni, da se funkcija ne bo izvajala kot običajna koda Python (v kateri se izvaja vsaka izvršljiva vrstica). Na primer, spremenimo našo prejšnjo funkcijo:

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

Rezultat tega je:

Only prints the first time!

O print() se izvede samo enkrat med korakom sledenja, to je takrat, ko se izvaja običajna koda Python. Naslednji klici funkcije izvajajo le operacije TenforFlow iz računskega grafa (operacije TensorFlow).

Vendar, če uporabljamo tf.print() namesto:

@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 vključuje tf.print() v svojem računskem grafu, saj gre za operacijo TensorFlow – ne za običajno funkcijo Python.

Najboljše prakse z @tf.function

Pisanje kode z operacijami TensorFlow

Kot smo pravkar pokazali, računski graf prezre nekatere dele kode. Zaradi tega je težko napovedati obnašanje funkcije pri kodiranju z »normalno« kodo Python, kot smo pravkar videli pri print(). Bolje je, da svojo funkcijo kodirate z operacijami TensorFlow, kadar je to primerno, da se izognete nepričakovanemu vedenju.

Na primer, for in while zanke se lahko pretvorijo v enakovredno zanko TensorFlow ali pa tudi ne. Zato je bolje, če je mogoče, zanko »za« zapisati kot vektorizirano operacijo. To bo izboljšalo delovanje vaše kode in zagotovilo pravilno sledenje vaši funkciji.

Kot primer razmislite o naslednjem:

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)

Koda z operacijami TensorFlow je bistveno hitrejša.

Izogibajte se sklicevanju na globalne spremenljivke

Razmislite o naslednji kodi:

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)

Prvič okrašena funkcija power() bila klicana, je bila izhodna vrednost pričakovana 4. Drugič pa je funkcija prezrla vrednost y je bil spremenjen. To se zgodi, ker je vrednost globalnih spremenljivk Python zamrznjena za funkcijo po sledenju.

Boljši način bi bila uporaba tf.Variable() za vse vaše spremenljivke in oboje posredujte kot argumenta vaši funkciji.

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)

Razhroščevanje [e-pošta zaščitena]_s

Na splošno želite razhroščevati svojo funkcijo v željnem načinu in jih nato okrasiti z @tf.function potem, ko se vaša koda izvaja pravilno, ker so sporočila o napakah v eager načinu bolj informativna.

Nekatere pogoste težave so tipske napake in napake oblike. Napake tipa se zgodijo, ko pride do neujemanja v vrsti spremenljivk, vključenih v operacijo:

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]

Tipske napake se zlahka prikradejo in jih je mogoče zlahka popraviti s pretvorbo spremenljivke v drug tip:

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

Napake oblike se zgodijo, ko vaši tenzorji nimajo oblike, ki jo zahteva vaša operacija:

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]

Eno priročno orodje za popravljanje obeh vrst napak je interaktivni razhroščevalnik Python, ki ga lahko samodejno pokličete v Jupyter Notebook z %pdb. S tem lahko kodirate svojo funkcijo in jo izvajate skozi nekaj običajnih primerov uporabe. Če pride do napake, se odpre interaktivni poziv. Ta poziv vam omogoča, da greste gor in dol po slojih abstrakcije v vaši kodi in preverite vrednosti, tipe in oblike vaših spremenljivk TensorFlow.

zaključek

Videli smo, kako je TensorFlow tf.function() naredi vaše delovanje učinkovitejše in kako @tf.function dekorater uporabi funkcijo za vašo.

Ta pospešitev je uporabna pri funkcijah, ki bodo večkrat klicane, kot so koraki usposabljanja po meri za modele strojnega učenja.