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

סקירה כללית

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

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

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

שרטוט

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

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

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

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

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

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

האתגרים הכרוכים בפעולות fused

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

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

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

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

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

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

דוגמה:

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 Translationer

מטמיעים את הפעולה המשולבת כפעולה בהתאמה אישית ב-TFLite – ראו instructions.

שימו לב שהשם שבאמצעותו צריך לרשום את ההפעלה צריך להיות דומה לשם שמצוין במאפיין 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 לפעולות משולבות של TensorFlow Lite:

שרטוט

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

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

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

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

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

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

המרה ל-TensorFlow Lite

משתמשים ב-API של TFLiteConverter.from_saved_model כדי לבצע המרה ל-TensorFlow Lite.

אפשרויות מתקדמות

עכשיו מתוארים פרטים כלליים על התכנון הכללי בהמרה לפעולות משולבות ב-TensorFlow Lite.

כתיבת פעולות ב-TensorFlow

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

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

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

  @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 כפי שמתואר למעלה, מאפשרת לפתח תשתית כללית כדי לזהות פעולות כאלה ולהמיר אותן לפעולות משולבות של TensorFlow Lite.

הרחבת הממיר של TensorFlow Lite

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

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

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

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

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

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

נבחן בפירוט את המעבר ל-Fusion. הכרטיס הזה מבצע את הפעולות הבאות:

  1. עוברים על כל הפונקציות במודול MLIR.
  2. אם לפונקציה יש את המאפיין tf._embeds על סמך ערך המאפיין, היא מפעילה את שירות הפעולות המתאים של Fusion.
  3. כלי העזר של הפעולה Fusion פועל על האופרנדים והמאפיינים של הפונקציה (שמשמשים כממשק להמרה) ומחליף את גוף הפונקציה בגוף פונקציה מקביל שמכיל את הפעולה המשולבת.
  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 */
}

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

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