הבנת ה-@tf.function Decorator של TensorFlow

מבוא

שיפור הביצועים של לולאת אימון יכול לחסוך שעות של זמן מחשוב בעת אימון מודלים של למידת מכונה. אחת הדרכים לשיפור הביצועים של קוד TensorFlow היא שימוש ב- tf.function() decorator - שינוי פשוט בשורה אחת שיכול לגרום לפונקציות שלך לפעול מהר יותר באופן משמעותי.

במדריך קצר זה נסביר כיצד 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)

עם זאת, אם נעביר את הפונקציה היקרה ל-a 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 בשני מצבים: מצב להוט ו מצב גרף. מצב להוט הוא הדרך הסטנדרטית והאינטראקטיבית להפעיל קוד: בכל פעם שאתה קורא לפונקציה, היא מבוצעת.

מצב גרף, לעומת זאת, הוא קצת שונה. במצב גרף, לפני ביצוע הפונקציה, 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 המקבילה. לכן, עדיף לכתוב את לולאת ה-"for" שלך כפעולה וקטורית, אם אפשר. זה ישפר את הביצועים של הקוד שלך ויבטיח שהפונקציה שלך תתחקה כראוי.

לדוגמה, שקול את הדברים הבאים:

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 [מוגן בדוא"ל]_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 מעצב מחיל את הפונקציה על שלך.

האצה זו שימושית בפונקציות שייקראו פעמים רבות, כגון שלבי אימון מותאמים אישית עבור מודלים של למידת מכונה.