Zrozumienie dekoratora @tf.function TensorFlow

Wprowadzenie

Poprawa wydajności pętli szkoleniowej może zaoszczędzić godziny czasu obliczeniowego podczas uczenia modeli uczenia maszynowego. Jednym ze sposobów poprawy wydajności kodu TensorFlow jest użycie tf.function() dekorator – prosta, jednowierszowa zmiana, która może znacznie przyspieszyć działanie Twoich funkcji.

W tym krótkim przewodniku wyjaśnimy, w jaki sposób tf.function() poprawia wydajność i zapoznaj się z najlepszymi praktykami.

Dekoratory Pythona i tf.funkcja()

W Pythonie dekorator to funkcja, która modyfikuje zachowanie innych funkcji. Załóżmy na przykład, że w komórce notatnika wywołujesz następującą funkcję:

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)

Jeśli jednak przekażemy kosztowną funkcję do a tf.function():

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

Dostajemy quicker_computation() – nowa funkcja, która działa znacznie szybciej niż poprzednia:

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

Więc, tf.function() modyfikuje some_costly_computation() i wyprowadza quicker_computation() funkcjonować. Dekoratorzy modyfikują również funkcje, więc naturalnym było wykonanie tf.function() również dekorator.

Korzystanie z notacji dekoratorskiej jest takie samo jak wywoływanie 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)

Jak działa tf.function() Praca?

Jak to możliwe, że niektóre funkcje działają 2-3 razy szybciej?

Kod TensorFlow można uruchomić w dwóch trybach: tryb chętny i tryb wykresu. Tryb Eager to standardowy, interaktywny sposób uruchamiania kodu: za każdym razem, gdy wywołujesz funkcję, jest ona wykonywana.

Tryb wykresu jest jednak nieco inny. W trybie wykresu, przed wykonaniem funkcji, TensorFlow tworzy wykres obliczeń, który jest strukturą danych zawierającą operacje wymagane do wykonania funkcji. Wykres obliczeń pozwala TensorFlow na uproszczenie obliczeń i znalezienie możliwości zrównoleglenia. Wykres izoluje również funkcję od nadrzędnego kodu Pythona, umożliwiając jej wydajne uruchamianie na wielu różnych urządzeniach.

Funkcja ozdobiona @tf.function odbywa się w dwóch krokach:

1. W pierwszym kroku TensorFlow wykonuje kod Pythona dla funkcji i kompiluje wykres obliczeń, opóźniając wykonanie dowolnej operacji TensorFlow.
2. Następnie uruchamiany jest wykres obliczeń.

Pierwszy krok zostanie pominięty, jeśli nie będzie potrzeby tworzenia nowego wykresu obliczeniowego. Poprawia to wydajność funkcji, ale oznacza również, że funkcja nie będzie działać jak zwykły kod Pythona (w którym wykonywany jest każdy wiersz wykonywalny). Na przykład zmodyfikujmy naszą poprzednią funkcję:

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

To skutkuje:

Only prints the first time!

Połączenia print() jest wykonywany tylko raz podczas kroku śledzenia, czyli gdy uruchamiany jest zwykły kod Pythona. Następne wywołania funkcji wykonują tylko operacje TenforFlow z wykresu obliczeń (operacje TensorFlow).

Jeśli jednak skorzystamy tf.print() zamiast:

@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 zawiera tf.print() w swoim grafie obliczeniowym, ponieważ jest to operacja TensorFlow – a nie zwykła funkcja Pythona.

Najlepsze praktyki z @tf.function

Pisanie kodu za pomocą TensorFlow Operations

Jak właśnie pokazaliśmy, niektóre części kodu są ignorowane przez wykres obliczeń. To sprawia, że ​​trudno jest przewidzieć zachowanie funkcji podczas kodowania „normalnym” kodem Pythona, jak właśnie widzieliśmy z print(). Lepiej jest zakodować swoją funkcję za pomocą operacji TensorFlow, gdy ma to zastosowanie, aby uniknąć nieoczekiwanego zachowania.

Na przykład, for i while pętle mogą, ale nie muszą, zostać przekształcone w równoważną pętlę TensorFlow. Dlatego lepiej jest napisać pętlę „for” jako operację wektorową, jeśli to możliwe. Poprawi to wydajność kodu i zapewni prawidłowe śledzenie funkcji.

Jako przykład rozważ następujące kwestie:

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)

Kod z operacjami TensorFlow jest znacznie szybszy.

Unikaj odniesień do zmiennych globalnych

Rozważ następujący kod:

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)

Po raz pierwszy urządzona funkcja power() została wywołana, wartość wyjściowa była oczekiwana 4. Jednak za drugim razem funkcja zignorowała, że ​​​​wartość y został zmieniony. Dzieje się tak, ponieważ wartość zmiennych globalnych Pythona jest zamrożona dla funkcji po śledzeniu.

Lepszym sposobem byłoby użycie tf.Variable() dla wszystkich zmiennych i przekaż oba jako argumenty do swojej funkcji.

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)

Debugowanie [email chroniony]_s

Ogólnie rzecz biorąc, chcesz debugować swoją funkcję w trybie przyspieszonym, a następnie udekorować je za pomocą @tf.function po poprawnym uruchomieniu kodu, ponieważ komunikaty o błędach w trybie przyspieszonym zawierają więcej informacji.

Niektóre typowe problemy to błędy typu i błędy kształtu. Błędy typu występują, gdy występuje niezgodność typu zmiennych biorących udział w operacji:

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]

Błędy typograficzne łatwo się wkradają i można je łatwo naprawić, rzutując zmienną na inny typ:

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

Błędy kształtu pojawiają się, gdy twoje tensory nie mają kształtu wymaganego przez twoją operację:

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]

Jednym z wygodnych narzędzi do naprawy obu rodzajów błędów jest interaktywny debuger Pythona, który można wywołać automatycznie w notatniku Jupyter za pomocą %pdb. Korzystając z tego, możesz zakodować swoją funkcję i uruchomić ją w kilku typowych przypadkach użycia. Jeśli wystąpi błąd, otwiera się interaktywny monit. Ten monit umożliwia poruszanie się w górę i w dół warstw abstrakcji w kodzie i sprawdzanie wartości, typów i kształtów zmiennych TensorFlow.

Wnioski

Widzieliśmy, jak działa TensorFlow tf.function() sprawia, że ​​Twoja funkcja jest bardziej wydajna i jak @tf.function dekorator stosuje tę funkcję do Twojej.

Przyspieszenie to jest przydatne w funkcjach, które będą wywoływane wielokrotnie, takich jak niestandardowe kroki uczenia dla modeli uczenia maszynowego.