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

סקירה כללית

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

המרת הפעולה המורכבת ל-tf.function

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

דוגמה,

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

יש להטמיע את הפעולה המשולבת כפעולה בהתאמה אישית של TFLite - מידע נוסף הוראות להגדרה.

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

דוגמה לאופציה בדוגמה היא

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

בכרטיס להכין פונקציות מורכבות, מוסיפים את הפלאגין המרות .

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

המרה ל-LiteRT

משתמשים ב TFLiteConverter.from_saved_model API להמרה ל-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 שהושק מוקדם יותר השנה נתמך רק לייבא מודלים של TensorFlow כתרשים שבו כל המשתנים הוחלפו את הערכים הקבועים התואמים. זה לא עובד עבור היתוך פעולה בתרשימים כאלה כל הפונקציות מסומנות כך שאפשר להפוך את המשתנים קבועים.

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

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

1. ייבוא של מודלים שמורים מ-TensorFlow ל-MLIR
2. מיזוג פעולות מורכבות
3. ניתוח של יכולת שינוי משתנה
4. להקפיא את כל המשתנים לקריאה בלבד

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

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

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

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

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

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());
  }