Κατανόηση του @tf.function Decorator του TensorFlow

Εισαγωγή

Η βελτίωση της απόδοσης ενός βρόχου εκπαίδευσης μπορεί να εξοικονομήσει ώρες υπολογιστικού χρόνου κατά την εκπαίδευση μοντέλων μηχανικής εκμάθησης. Ένας από τους τρόπους βελτίωσης της απόδοσης του κώδικα TensorFlow είναι η χρήση του tf.function() διακοσμητής – μια απλή αλλαγή μιας γραμμής που μπορεί να κάνει τις λειτουργίες σας να εκτελούνται σημαντικά πιο γρήγορα.

Σε αυτόν τον σύντομο οδηγό, θα εξηγήσουμε πώς tf.function() βελτιώνει την απόδοση και ρίξτε μια ματιά σε ορισμένες βέλτιστες πρακτικές.

Python Decorators και tf.function()

Στην Python, ο διακοσμητής είναι μια συνάρτηση που τροποποιεί τη συμπεριφορά άλλων συναρτήσεων. Για παράδειγμα, ας υποθέσουμε ότι καλείτε την ακόλουθη συνάρτηση σε ένα κελί φορητού υπολογιστή:

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)

Ωστόσο, αν περάσουμε την δαπανηρή συνάρτηση σε α tf.function():

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

Παίρνουμε quicker_computation() – μια νέα λειτουργία που εκτελεί πολύ πιο γρήγορα από την προηγούμενη:

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

Έτσι tf.function() τροποποιεί some_costly_computation() και βγάζει το quicker_computation() λειτουργία. Οι διακοσμητές τροποποιούν επίσης λειτουργίες, οπότε ήταν φυσικό να γίνει tf.function() επίσης ένας διακοσμητής.

Η χρήση της σημειογραφίας διακοσμητή είναι ίδια με την κλήση 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)

Πώς κάνει tf.function() Εργασία?

Πώς γίνεται να κάνουμε ορισμένες λειτουργίες να εκτελούνται 2-3 φορές πιο γρήγορα;

Ο κώδικας TensorFlow μπορεί να εκτελεστεί σε δύο τρόπους: ανυπόμονη λειτουργία και λειτουργία γραφήματος. Η λειτουργία Eager είναι ο τυπικός, διαδραστικός τρόπος εκτέλεσης κώδικα: κάθε φορά που καλείτε μια συνάρτηση, αυτή εκτελείται.

Η λειτουργία γραφήματος, ωστόσο, είναι λίγο διαφορετική. Στη λειτουργία γραφήματος, πριν από την εκτέλεση της συνάρτησης, το TensorFlow δημιουργεί ένα υπολογιστικό γράφημα, το οποίο είναι μια δομή δεδομένων που περιέχει τις λειτουργίες που απαιτούνται για την εκτέλεση της συνάρτησης. Το γράφημα υπολογισμού επιτρέπει στο TensorFlow να απλοποιεί τους υπολογισμούς και να βρίσκει ευκαιρίες για παραλληλοποίηση. Το γράφημα απομονώνει επίσης τη συνάρτηση από τον υπερκείμενο κώδικα Python, επιτρέποντάς της να εκτελείται αποτελεσματικά σε πολλές διαφορετικές συσκευές.

Μια λειτουργία διακοσμημένη με @tf.function εκτελείται σε δύο βήματα:

1. Στο πρώτο βήμα, το TensorFlow εκτελεί τον κώδικα Python για τη συνάρτηση και μεταγλωττίζει ένα υπολογιστικό γράφημα, καθυστερώντας την εκτέλεση οποιασδήποτε λειτουργίας TensorFlow.
2. Στη συνέχεια, εκτελείται το γράφημα υπολογισμού.

Το πρώτο βήμα θα παραλειφθεί εάν δεν χρειάζεται να δημιουργηθεί ένα νέο υπολογιστικό γράφημα. Αυτό βελτιώνει την απόδοση της συνάρτησης αλλά σημαίνει επίσης ότι η συνάρτηση δεν θα εκτελείται όπως ο κανονικός κώδικας Python (στον οποίο εκτελείται κάθε εκτελέσιμη γραμμή). Για παράδειγμα, ας τροποποιήσουμε την προηγούμενη συνάρτησή μας:

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

Αυτο εχει ως αποτελεσμα:

Only prints the first time!

Η print() εκτελείται μόνο μία φορά κατά τη διάρκεια του βήματος ανίχνευσης, δηλαδή όταν εκτελείται ο κανονικός κώδικας Python. Οι επόμενες κλήσεις στη συνάρτηση εκτελούν μόνο λειτουργίες TenforFlow από το γράφημα υπολογισμού (πράξεις TensorFlow).

Ωστόσο, αν χρησιμοποιήσουμε tf.print() αντι αυτου:

@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 περιλαμβάνει tf.print() στο υπολογιστικό του γράφημα καθώς είναι μια λειτουργία TensorFlow – όχι μια κανονική συνάρτηση Python.

Βέλτιστες πρακτικές με @tf.function

Σύνταξη κώδικα με λειτουργίες TensorFlow

Όπως μόλις δείξαμε, ορισμένα τμήματα του κώδικα αγνοούνται από το γράφημα υπολογισμού. Αυτό καθιστά δύσκολη την πρόβλεψη της συμπεριφοράς της συνάρτησης κατά την κωδικοποίηση με "κανονικό" κώδικα Python, όπως μόλις είδαμε με print(). Είναι προτιμότερο να κωδικοποιήσετε τη συνάρτησή σας με λειτουργίες TensorFlow όταν αυτό ισχύει για να αποφύγετε απροσδόκητη συμπεριφορά.

Για παράδειγμα, for και while Οι βρόχοι μπορεί να μετατραπούν ή να μην μετατραπούν στον ισοδύναμο βρόχο TensorFlow. Επομένως, είναι καλύτερο να γράψετε τον βρόχο «για» ως διανυσματική πράξη, αν είναι δυνατόν. Αυτό θα βελτιώσει την απόδοση του κώδικά σας και θα διασφαλίσει ότι η λειτουργία σας παρακολουθεί σωστά.

Ως παράδειγμα, λάβετε υπόψη τα ακόλουθα:

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)

Ο κώδικας με τις λειτουργίες TensorFlow είναι πολύ πιο γρήγορος.

Αποφύγετε τις αναφορές σε καθολικές μεταβλητές

Σκεφτείτε τον ακόλουθο κώδικα:

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)

Η πρώτη φορά η διακοσμημένη λειτουργία power() κλήθηκε, η τιμή εξόδου ήταν η αναμενόμενη 4. Ωστόσο, τη δεύτερη φορά, η συνάρτηση αγνόησε ότι η τιμή του y άλλαξε. Αυτό συμβαίνει επειδή η τιμή των καθολικών μεταβλητών Python έχει παγώσει για τη συνάρτηση μετά τον εντοπισμό.

Ένας καλύτερος τρόπος θα ήταν να χρησιμοποιήσετε tf.Variable() για όλες τις μεταβλητές σας και περάστε και τις δύο ως ορίσματα στη συνάρτησή σας.

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 [προστασία μέσω email]_s

Γενικά, θέλετε να διορθώσετε τη λειτουργία σας σε κατάσταση ανυπομονησίας και, στη συνέχεια, να τη διακοσμήσετε @tf.function αφού ο κώδικάς σας εκτελείται σωστά, επειδή τα μηνύματα σφάλματος στη λειτουργία ανυπομονησίας είναι πιο ενημερωτικά.

Μερικά κοινά προβλήματα είναι τα σφάλματα τύπου και τα σφάλματα σχήματος. Τα σφάλματα τύπου συμβαίνουν όταν υπάρχει αναντιστοιχία στον τύπο των μεταβλητών που εμπλέκονται σε μια λειτουργία:

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]

Τα σφάλματα τύπου εισχωρούν εύκολα και μπορούν εύκολα να διορθωθούν μεταφέροντας μια μεταβλητή σε διαφορετικό τύπο:

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

Τα σφάλματα σχήματος συμβαίνουν όταν οι τανυστές σας δεν έχουν το σχήμα που απαιτεί η λειτουργία σας:

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]

Ένα βολικό εργαλείο για τη διόρθωση και των δύο ειδών σφαλμάτων είναι το διαδραστικό πρόγραμμα εντοπισμού σφαλμάτων Python, το οποίο μπορείτε να καλέσετε αυτόματα σε ένα σημειωματάριο Jupyter χρησιμοποιώντας %pdb. Χρησιμοποιώντας αυτό, μπορείτε να κωδικοποιήσετε τη συνάρτησή σας και να την εκτελέσετε μέσω ορισμένων συνηθισμένων περιπτώσεων χρήσης. Εάν υπάρχει σφάλμα, ανοίγει μια διαδραστική προτροπή. Αυτή η προτροπή σάς επιτρέπει να ανεβοκατεβείτε τα επίπεδα αφαίρεσης στον κώδικά σας και να ελέγξετε τις τιμές, τους τύπους και τα σχήματα των μεταβλητών σας TensorFlow.

Συμπέρασμα

Είδαμε πώς είναι το TensorFlow tf.function() κάνει τη λειτουργία σας πιο αποτελεσματική και πώς το @tf.function ο διακοσμητής εφαρμόζει τη λειτουργία στο δικό σας.

Αυτή η επιτάχυνση είναι χρήσιμη σε λειτουργίες που θα καλούνται πολλές φορές, όπως προσαρμοσμένα βήματα εκπαίδευσης για μοντέλα μηχανικής εκμάθησης.