המרת RNN של TensorFlow ל-LiteRT

סקירה כללית

‫LiteRT תומך בהמרת מודלים של TensorFlow RNN לפעולות של LiteRT fused LSTM. הפעולות הממוזגות נועדו למקסם את הביצועים של ההטמעות הבסיסיות שלהן בקרנל, וגם לספק ממשק ברמה גבוהה יותר להגדרת טרנספורמציות מורכבות כמו קוונטיזציה.

מכיוון שיש הרבה וריאציות של ממשקי RNN API ב-TensorFlow, הגישה שלנו היא כפולה:

  1. תמיכה מקורית בממשקי API סטנדרטיים של TensorFlow RNN, כמו Keras LSTM. זו האפשרות המומלצת.
  2. מספק ממשק לתוך תשתית ההמרות עבור הטמעות של רשתות RNN שהוגדרו על ידי המשתמש כדי לחבר אותן ולהמיר אותן ל-LiteRT. אנחנו מספקים כמה דוגמאות מוכנות לשימוש להמרה כזו באמצעות ממשקי ה-RNN של LSTMCellSimple ו-LayerNormalizedLSTMCellSimple של lingvo.

Converter API

התכונה הזו היא חלק מגרסה TensorFlow 2.3. אפשר להשתמש בו גם דרך tf-nightly pip או מ-head.

הפונקציונליות הזו של המרות זמינה כשממירים ל-LiteRT באמצעות SavedModel או ישירות ממודל Keras. דוגמאות לשימוש

מהמודל שנשמר

# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)

# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

ממודל Keras

# build a Keras model
keras_model = build_keras_lstm(...)

# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

דוגמה

‫Keras LSTM to LiteRT Colab ממחיש את השימוש מקצה לקצה עם רכיב התרגום של LiteRT.

ממשקי API של TensorFlow RNNs נתמכים

אנחנו תומכים בהמרה מוכנה מראש של Keras LSTM ל-LiteRT. פרטים על אופן הפעולה זמינים בממשק Keras LSTM ובלוגיקת ההמרה כאן.

חשוב גם להדגיש את חוזה ה-LSTM של LiteRT ביחס להגדרת הפעולה של Keras:

  1. המאפיין 0 של טנסור הקלט הוא גודל האצווה.
  2. הממד 0 של טנזור recurrent_weight הוא מספר הפלטים.
  3. הטנסורים weight ו-recurrent_kernel מוחלפים.
  4. טנסורים של משקל מוזז, recurrent_kernel מוזז ו-bias מחולקים ל-4 טנסורים בגודל שווה לאורך המימד 0. הם תואמים ל-input gate,‏ forget gate,‏ cell ו-output gate.

וריאציות של Keras LSTM

הזמן הוא הממד העיקרי

המשתמשים יכולים לבחור אם להשתמש בפורמט time-major או לא. הפונקציה Keras LSTM מוסיפה מאפיין time-major במאפייני הגדרת הפונקציה. במקרה של LSTM חד-כיווני, אפשר פשוט למפות את המאפיין time major של unidirecional_sequence_lstm.

BiDirectional LSTM

אפשר להטמיע LSTM דו-כיווני באמצעות שתי שכבות Keras LSTM, אחת להעברה קדימה ואחת להעברה אחורה. דוגמאות זמינות כאן. אחרי שאנחנו רואים את המאפיין go_backward, אנחנו מזהים אותו כ-LSTM לאחור, ואז מקבצים את ה-LSTM קדימה ואחורה. העבודה הזו מתוכננת לעתיד. בשלב הזה, נוצרות שתי פעולות UnidirectionalSequenceLSTM במודל LiteRT.

דוגמאות להמרות LSTM שהוגדרו על ידי המשתמש

‫LiteRT מספק גם דרך להמיר יישומי LSTM שהוגדרו על ידי המשתמש. בדוגמה הבאה אנחנו משתמשים ב-LSTM של Lingvo כדי להמחיש איך אפשר ליישם את זה. פרטים נוספים זמינים כאן. בנוסף, אנחנו מספקים דוגמה להגדרה אחרת של LSTM ב-Lingvo בממשק lingvo.LayerNormalizedLSTMCellSimple ובלוגיקת ההמרה שלו כאן.

‫“Bring your own TensorFlow RNN” to LiteRT

אם ממשק ה-RNN של משתמש שונה מהממשקים הנתמכים הרגילים, יש כמה אפשרויות:

אפשרות 1: כותבים קוד מתאם ב-TensorFlow Python כדי להתאים את ממשק ה-RNN לממשק ה-RNN של Keras. כלומר, tf.function עם tf_implements annotation בפונקציה של ממשק ה-RNN שנוצר, זהה לזה שנוצר על ידי שכבת ה-LSTM של Keras. אחרי זה, אותו API להמרות שמשמש ל-Keras LSTM יפעל.

אפשרות 2: אם האפשרות הראשונה לא מתאימה (למשל, אם חסרה ב-Keras LSTM פונקציונליות מסוימת שמוצגת כרגע על ידי LiteRT fused LSTM op כמו layer normalization), אפשר להרחיב את LiteRT converter על ידי כתיבת קוד המרה מותאם אישית והוספתו אל prepare-composite-functions MLIR-pass כאן. צריך להתייחס לממשק של הפונקציה כמו אל חוזה API, והוא צריך לכלול את הארגומנטים שנדרשים להמרה לפעולות LiteRT LSTM משולבות – כלומר, קלט, הטיה, משקלים, הקרנה, נורמליזציה של השכבה וכו'. עדיף שהטנסורים שמועברים כארגומנטים לפונקציה הזו יהיו בעלי דרגה ידועה (כלומר, RankedTensorType ב-MLIR). כך קל יותר לכתוב קוד המרה שיכול להניח שהטנסורים האלה הם RankedTensorType, ולעזור להפוך אותם לטנסורים מדורגים שמתאימים לאופרנדים של אופרטור LiteRT משולב.

דוגמה מלאה לזרימת המרות כזו היא המרה של LSTMCellSimple של Lingvo ל-LiteRT.

ה-LSTMCellSimple ב-Lingvo מוגדר כאן. אפשר להמיר מודלים שאומנו באמצעות תא LSTM הזה ל-LiteRT באופן הבא:

  1. צריך לעטוף את כל השימושים ב-LSTMCellSimple ב-tf.function עם הערה מסוג tf_implements שמסומנת ככזו (לדוגמה, lingvo.LSTMCellSimple יהיה שם הערה טוב כאן). מוודאים שהפונקציה tf.function שנוצרת תואמת לממשק של הפונקציה שצפויה בקוד ההמרה. זהו חוזה בין מחבר המודל שמוסיף את ההערה לבין קוד ההמרה.

  2. הרחבת השלב prepare-composite-functions כדי להוסיף פעולת אופרטור מורכבת בהתאמה אישית להמרה של פעולת אופרטור LSTM משולבת ב-LiteRT. אפשר לעיין בקוד ההמרה של LSTMCellSimple.

    החוזה להמרה:

  3. טנסורים של משקל והטלה מוחלפים.

  4. הערכים {input, recurrent} עד {cell, input gate, forget gate, output gate} מחולצים על ידי חיתוך טנסור המשקלים המועבר.

  5. ההטיה {bias} ל{cell, input gate, forget gate, output gate} מחולצת על ידי חיתוך טנזור ההטיה.

  6. ההטלה מופקת על ידי חיתוך של טנזור ההטלה המועבר.

  7. המרת דומות נכתבת עבור LayerNormalizedLSTMCellSimple.

  8. אפשר לעשות שימוש חוזר בשאר התשתית של המרות LiteRT, כולל כל המעברים של MLIR שהוגדרו, וגם בייצוא הסופי ל-LiteRT flatbuffer.

בעיות או מגבלות ידועות

  1. בשלב הזה יש תמיכה רק בהמרת LSTM של Keras ללא מצב (התנהגות ברירת המחדל ב-Keras). המרת LSTM של Keras עם שמירת מצב היא עבודה עתידית.
  2. עדיין אפשר ליצור מודל של שכבת Keras LSTM עם שמירת מצב באמצעות שכבת Keras LSTM הבסיסית ללא שמירת מצב, ולנהל את המצב באופן מפורש בתוכנית המשתמש. עדיין אפשר להמיר תוכנית TensorFlow כזו ל-LiteRT באמצעות התכונה שמתוארת כאן.
  3. מודל LSTM דו-כיווני מוגדר כרגע כשתי פעולות UnidirectionalSequenceLSTM ב-LiteRT. הוא יוחלף בפעולה אחת של BidirectionalSequenceLSTM.