دمج عملية TensorFlow

نظرة عامة

توضّح هذه الصفحة التصميم والخطوات اللازمة لتحويل العمليات المركّبة في TensorFlow إلى عمليات مدمجة في LiteRT. هذه البنية الأساسية عامة الأغراض وتتيح تحويل أي عملية مركّبة في TensorFlow إلى عملية مدمجة مقابلة في LiteRT.

من الأمثلة على استخدام هذه البنية الأساسية دمج عمليات TensorFlow RNN في LiteRT، كما هو موضّح بالتفصيل هنا.

ما هي العمليات المدمجة؟

يمكن أن تكون عمليات TensorFlow عمليات أساسية، مثل tf.add، أو يمكن أن تكون مكوّنة من عمليات أساسية أخرى، مثل tf.einsum. تظهر العملية الأساسية كعُقدة واحدة في الرسم البياني لـ TensorFlow، بينما العملية المركّبة هي مجموعة من العُقد في الرسم البياني لـ TensorFlow. ويكون تنفيذ عملية مركّبة مكافئًا لتنفيذ كل عملية من العمليات الأساسية المكوّنة لها.

تتطابق العملية المدمجة مع عملية واحدة تشمل كل العمليات الحسابية التي يتم تنفيذها بواسطة كل عملية أساسية ضمن العملية المركّبة المقابلة.

مزايا العمليات المدمجة

تتوفّر العمليات المدمجة لتحقيق أقصى أداء ممكن لعمليات التنفيذ الأساسية في النواة، وذلك من خلال تحسين الحسابات الإجمالية وتقليل مساحة الذاكرة المستخدَمة. وهذا مفيد جدًا، خاصةً لأحمال عمل الاستدلال التي تتطلّب وقت استجابة منخفضًا والمنصات الجوّالة ذات الموارد المحدودة.

توفّر العمليات المدمجة أيضًا واجهة ذات مستوى أعلى لتحديد عمليات التحويل المعقّدة، مثل التكميم، والتي قد يكون من غير الممكن أو من الصعب جدًا تنفيذها على مستوى أكثر تفصيلاً.

تتضمّن LiteRT العديد من حالات العمليات المدمجة للأسباب الموضّحة أعلاه. تتطابق هذه العمليات المدمجة عادةً مع العمليات المركّبة في برنامج TensorFlow المصدر. تشمل أمثلة العمليات المركّبة في TensorFlow التي يتم تنفيذها كعملية مدمجة واحدة في LiteRT عمليات مختلفة لشبكات RNN، مثل LSTM التسلسلي أحادي الاتجاه وثنائي الاتجاه، والالتفاف (conv2d، وإضافة الانحياز، وrelu)، والاتصال الكامل (matmul، وإضافة الانحياز، و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 في توقيع التنفيذ.

مثال على عامل التشغيل في المثال هو

  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، أدخِل رمز الإحالة الناجحة.

في الاستخدامات الأكثر تقدّمًا، من الممكن تنفيذ عمليات تحويل معقّدة لمعاملات العملية المركّبة من أجل استخلاص معاملات العملية المدمجة. يمكنك الاطّلاع على رمز تحويل LSTM في Keras كمثال.

التحويل إلى LiteRT

استخدِم واجهة برمجة التطبيقات 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 الذي تم إصداره في وقت سابق من هذا العام يتيح فقط استيراد نماذج TensorFlow كرسم بياني مع استبدال جميع المتغيرات بقيمها الثابتة المقابلة. لا يعمل ذلك مع دمج العمليات لأنّ هذه الرسومات البيانية تتضمّن جميع الدوال المضمّنة حتى يمكن تحويل المتغيرات إلى ثوابت.

للاستفادة من tf.function مع ميزة experimental_implements أثناء عملية التحويل، يجب الحفاظ على الدوال حتى وقت لاحق من عملية التحويل.

نتيجةً لذلك، نفّذنا سير عمل جديدًا لاستيراد نماذج TensorFlow وتحويلها في المحوّل من أجل إتاحة حالة استخدام دمج العمليات المركّبة. على وجه التحديد، تشمل الميزات الجديدة التي تمت إضافتها ما يلي:

1. استيراد النماذج المحفوظة في TensorFlow إلى MLIR
2. دمج العمليات المركّبة
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());
  }