מיזוג אוויר של TensorFlow

סקירה כללית

בדף הזה מתואר העיצוב והשלבים הנדרשים להמרת פעולות מורכבות ב-TensorFlow לפעולות משולבות ב-LiteRT. התשתית הזו היא למטרות כלליות ותומכת בהמרה של כל פעולה מורכבת ב-TensorFlow לפעולה מקבילה ב-LiteRT.

דוגמה לשימוש בתשתית הזו היא מיזוג של פעולות TensorFlow RNN ל-LiteRT, כפי שמפורט כאן.

מהן פעולות משולבות

פעולות TensorFlow יכולות להיות פעולות פרימיטיביות, למשל tf.add, או שהן יכולות להיות מורכבות מפעולות פרימיטיביות אחרות, למשל tf.einsum. פעולה פרימיטיבית מוצגת כצומת יחיד בגרף TensorFlow, ואילו פעולה מורכבת היא אוסף של צמתים בגרף TensorFlow. ביצוע של פעולה מורכבת שווה לביצוע של כל אחת מהפעולות הפשוטות שמרכיבות אותה.

פעולה משולבת מתאימה לפעולה יחידה שכוללת את כל החישובים שמתבצעים על ידי כל פעולה פרימיטיבית בתוך הפעולה המורכבת המתאימה.

היתרונות של פעולות משולבות

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

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

ב-LiteRT יש הרבה מקרים של פעולות משולבות, מהסיבות שצוינו למעלה. בדרך כלל, הפעולות הממוזגות האלה תואמות לפעולות מורכבות בתוכנית TensorFlow המקורית. דוגמאות לפעולות מורכבות ב-TensorFlow שמיושמות כפעולה מאוחדת יחידה ב-LiteRT כוללות פעולות שונות של RNN כמו Unidirectional ו-Bidirectional sequence LSTM, קונבולוציה (conv2d, ‏ bias add, ‏ relu), חיבור מלא (matmul, ‏ bias add, ‏ relu) ועוד. ב-LiteRT, קוונטיזציה של LSTM מיושמת כרגע רק בפעולות של LSTM משולב.

אתגרים בפעולות משולבות

המרת פעולות מורכבות מ-TensorFlow לפעולות משולבות ב-LiteRT היא בעיה מורכבת. הסיבות לכך הן:

1. פעולות מורכבות מיוצגות בתרשים TensorFlow כקבוצה של פעולות פרימיטיביות ללא גבול מוגדר. יכול להיות מאוד קשה לזהות (למשל באמצעות התאמת תבניות) את תת-הגרף שמתאים לפעולה מורכבת כזו.

2. יכול להיות שיש יותר מהטמעה אחת של TensorFlow שמטרגטת פעולת LiteRT משולבת. לדוגמה, יש הרבה יישומי LSTM ב-TensorFlow (Keras,‏ Babelfish/lingvo וכו'), וכל אחד מהם מורכב מפעולות פרימיטיביות שונות, אבל עדיין אפשר להמיר את כולם לאותה פעולת LSTM משולבת ב-LiteRT.

לכן, ההמרה של פעולות משולבות היא די מאתגרת.

המרת פעולה מורכבת לפעולה מותאמת אישית של TFLite (מומלץ)

עוטפים את הפעולה המורכבת בתג tf.function

במקרים רבים, אפשר למפות חלק מהמודל לפעולה יחידה ב-TFLite. כך אפשר לשפר את הביצועים כשכותבים הטמעה אופטימלית לפעולות ספציפיות. כדי ליצור פעולה משולבת ב-TFLite, צריך לזהות את החלק בתרשים שמייצג פעולה משולבת ולעטוף אותו ב-tf.function עם המאפיין experimental_implements ל-tf.function, שכולל את ערך המאפיין tfl_fusable_op עם הערך true. אם הפעולה המותאמת אישית מקבלת מאפיינים, צריך להעביר אותם כחלק מאותו experimental_implements.

לדוגמה,

def get_implements_signature():
  implements_signature = [
    # 'name' will be used as a name for the operation.
    'name: "my_custom_fused_op"',
    # attr "tfl_fusable_op" is required to be set with true value.
    'attr {key: "tfl_fusable_op" value { b: true } }',
    # Example attribute "example_option" that the op accepts.
    'attr {key: "example_option" value { i: %d } }' % 10
  ]
  return ' '.join(implements_signature)

@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
  # An empty function that represents pre/post processing example that
  # is not represented as part of the Tensorflow graph.
  output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
  output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
  return output_1, output_2

class TestModel(tf.Module):
  def __init__(self):
    super(TestModel, self).__init__()
    self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
    self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
  ])
  def simple_eval(self, input_a, input_b):
    return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))

הערה: אין צורך להגדיר את allow_custom_ops בממיר כי מאפיין tfl_fusable_op מרמז על כך.

הטמעה של פעולה מותאמת אישית ורישום שלה ב-TFLite Interpreter

מטמיעים את הפעולה הממוזגת כפעולה מותאמת אישית של TFLite – הוראות

שימו לב: השם שבו רוצים לרשום את האופרטור צריך להיות דומה לשם שצוין במאפיין name בחתימת ההטמעה.

דוגמה ל-op בדוגמה היא

  TfLiteRegistration reg = {};
  // This name must match the name specified in the implements signature.
  static constexpr char kOpName[] = "my_custom_fused_op";
  reg.custom_name = kOpName;
  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

המרת פעולה מורכבת לפעולה משולבת (מתקדם)

הארכיטקטורה הכוללת להמרת פעולות מורכבות של TensorFlow לפעולות משולבות של LiteRT מוצגת בהמשך:

עוטפים את הפעולה המורכבת בתג tf.function

בקוד המקור של מודל TensorFlow, מזהים את הפעולה המורכבת ומבצעים לה הפשטה ל-tf.function עם הערת הפונקציה experimental_implements. דוגמה להטמעת חיפוש הפונקציה מגדירה את הממשק, והארגומנטים שלה צריכים לשמש להטמעה של לוגיקת ההמרה.

כתיבת קוד המרה

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

בשלב prepare-composite-functions, מחברים את קוד ההמרה.

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

המרת LiteRT

משתמשים ב-API‏ TFLiteConverter.from_saved_model כדי להמיר ל-LiteRT.

הפרטים הטכניים

עכשיו נתאר פרטים ברמה גבוהה של העיצוב הכולל של המעבר לפעולות משולבות ב-LiteRT.

הרכבת פעולות ב-TensorFlow

השימוש ב-tf.function עם מאפיין הפונקציה experimental_implements מאפשר למשתמשים ליצור באופן מפורש פעולות חדשות באמצעות פעולות פרימיטיביות של TensorFlow, ולציין את הממשק שהפעולה המורכבת שמתקבלת מטמיעה. האפשרות הזו שימושית מאוד כי היא מספקת:

1. גבול מוגדר היטב לפעולה המורכבת בתרשים TensorFlow הבסיסי.
2. מציינים במפורש את הממשק שהפעולה הזו מיישמת. הארגומנטים של tf.function תואמים לארגומנטים של הממשק הזה.

לדוגמה, נניח שיש פעולה מורכבת שמוגדרת להטמעת חיפוש הטמעה. המיפוי הזה מתבצע לפעולה משולבת ב-LiteRT.

  @tf.function(
        experimental_implements="embedding_lookup")
    def EmbFprop(embs, ids_vec):
      """Embedding forward prop.

      Effectively, it computes:
        num = size of ids_vec
        rets = zeros([num, embedding dim])
        for i in range(num):
          rets[i, :] = embs[ids_vec[i], :]
        return rets

      Args:
        embs: The embedding matrix.
        ids_vec: A vector of int32 embedding ids.

      Returns:
        The result of embedding lookups. A matrix of shape
        [num ids in ids_vec, embedding dims].
      """
      num = tf.shape(ids_vec)[0]
      rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))

      def EmbFpropLoop(i, embs, ids_vec, rets):
        # row_id = ids_vec[i]
        row_id = tf.gather(ids_vec, i)
        # row = embs[row_id]
        row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
        # rets[i] = row
        rets = inplace_ops.alias_inplace_update(rets, [i], row)
        return embs, ids_vec, rets

      _, _, rets = functional_ops.For(
          start=0,
          limit=num,
          delta=1,
          inputs=[embs, ids_vec, rets],
          body=EmbFpropLoop,
          rewrite_with_while=compiled)
      if len(weight_shape) > 2:
        rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
      return rets

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

הרחבת הממיר LiteRT

הכלי LiteRT Converter שפורסם מוקדם יותר השנה תמך רק בייבוא של מודלים של TensorFlow כגרף, כשכל המשתנים מוחלפים בערכים הקבועים התואמים שלהם. השיטה הזו לא מתאימה למיזוג פעולות, כי בכל הגרפים האלה הפונקציות מוטמעות כך שאפשר להפוך את המשתנים לקבועים.

כדי להשתמש ב-tf.function עם התכונה experimental_implements במהלך תהליך ההמרה, צריך לשמור את הפונקציות עד לשלב מאוחר יותר בתהליך ההמרה.

לכן, הטמענו תהליך עבודה חדש לייבוא ולהמרה של מודלים של TensorFlow בממיר, כדי לתמוך בתרחיש השימוש של מיזוג פעולות מורכבות. התכונות החדשות שנוספו הן:

1. ייבוא של מודלים שמורים של TensorFlow אל MLIR
2. fuse composite operations
3. ניתוח של שינוי משתנים
4. הקפאה של כל המשתנים לקריאה בלבד

כך אנחנו יכולים לבצע מיזוג של פעולות באמצעות הפונקציות שמייצגות את הפעולות המורכבות לפני הטמעת הפונקציות והקפאת המשתנים.

הטמעה של מיזוג פעולות

בואו נסתכל על שלב מיזוג הפעולות בפירוט. הכרטיס הזה:

1. לולאה בכל הפונקציות במודול MLIR.
2. אם לפונקציה יש מאפיין tf._implements, הפונקציה קוראת לכלי המתאים לאיחוד פעולות על סמך ערך המאפיין.
3. כלי המיזוג של הפעולות פועל על האופרנדים והמאפיינים של הפונקציה (שמשמשים כממשק להמרה), ומחליף את גוף הפונקציה בגוף פונקציה שווה ערך שמכיל את הפעולה הממוזגת.
4. במקרים רבים, הגוף שהוחלף יכיל פעולות אחרות מלבד הפעולה הממוזגת. הן תואמות לחלק מהטרנספורמציות הסטטיות באופרנדים של הפונקציה, כדי לקבל את האופרנדים של הפעולה הממוזגת. מכיוון שכל החישובים האלה יכולים להתבצע מראש, הם לא יופיעו ב-flatbuffer המיוצא, שבו תופיע רק הפעולה המאוחדת.

קטע הקוד הבא מתוך הכרטיס מציג את תהליך העבודה העיקרי:

void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
                                                        StringAttr attr) {
  if (attr.getValue() == "embedding_lookup") {
    func.eraseBody();
    func.addEntryBlock();
    // Convert the composite embedding_lookup function body to a
    // TFLite fused embedding_lookup op.
    ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
    if (failed(convert_embedded_lookup.VerifySignature())) {
      return signalPassFailure();
    }
    convert_embedded_lookup.RewriteFunc();
  } else if (attr.getValue() == mlir::TFL::kKerasLstm) {
     func.eraseBody();
     func.addEntryBlock();
     OpBuilder builder(func.getBody());
     if (failed(ConvertKerasLSTMLayer(func, &builder))) {
       return signalPassFailure();
     }
  } else if (.....) /* Other fusions can plug in here */
}

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

void RewriteFunc() {
    Value lookup = func_.getArgument(1);
    Value value = func_.getArgument(0);
    auto output_type = func_.getType().getResult(0);

    OpBuilder builder(func_.getBody());
    auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
        func_.getLoc(), output_type, lookup, value);

    builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
  }