היתוך פעולת TensorFlow

סקירה כללית

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

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

מהן פעולות התמזגות

צִיוּר

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

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

היתרונות של פעולות מתמזגות

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

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

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

אתגרים עם פעולות מתמזגות

המרת פעולות מורכבות מ-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_implements" ל- tf.function , שיש לה תכונה ערך tfl_fusable_op עם ערך true . אם הפעולה המותאמת אישית לוקחת תכונות אז העבר אותן כחלק מאותם "יישמים_ניסויים".

דוגמא,

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 Custom - ראה הוראות .

שים לב שהשם לרשום את האופ צריך להיות דומה לשם המצוין בתכונת 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);

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

הארכיטקטורה הכוללת להמרת פעולות מרוכבות של TensorFlow לפעולות ממוזגות של TensorFlow Lite היא להלן:

צִיוּר

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

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

כתוב קוד המרה

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

במעבר ה-preparation-composite-functions, תוסף בקוד ההמרה שלך.

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

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

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

יישום היתוך פעולה

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

  1. עברו בלולאה בין כל הפונקציות במודול MLIR.
  2. אם לפונקציה יש את התכונה tf._implements, בהתבסס על ערך התכונה, מתקשרת לכלי השירות המתאים לפעולה fusion.
  3. כלי ה-operation 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());
  }