دمج عملية TensorFlow

نظرة عامة

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

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

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

رسم

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

تتجاوب العملية المدمَجة مع عملية واحدة تشكِّل جميع العملية الحسابية التي تنفذها كل عملية أساسية داخل العملية المركّبة.

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

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

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

توجد في LiteRT العديد من مثيلات العمليات المدمجة للأسباب تمامًا أعلاه. وعادةً ما تتوافق هذه العمليات المدمَجة مع المركّب العمليات في برنامج TensorFlow المصدر. أمثلة على العمليات المركّبة في TensorFlow التي يتم تنفيذها كعملية واحدة مدمجة في LiteRT تضمين العديد من عمليات RNN، مثل المتتابعة الأحادية الاتجاه والتسلسل ثنائي الاتجاه LSTM، الالتفاف (conv2d، الانحياز إضافة، relu)، متصل بالكامل (matmul، انحياز add، relu) وغير ذلك. في LiteRT، يتم حاليًا قياس LSTM فقط تنفيذها في عمليات LSTM المدمجة.

تحديات دمج العمليات

تحويل العمليات المركّبة من TensorFlow إلى العمليات المدمجة في LiteRT مشكلة صعبة. ويرجع ذلك إلى ما يلي:

  1. يتم تمثيل العمليات المركّبة في الرسم البياني TensorFlow كمجموعة من العمليات الأولية بدون حدود محددة جيدًا. من الممكن أن يسير يصعب تحديد (على سبيل المثال عن طريق مطابقة الأنماط) الرسم البياني الفرعي المقابلة لمثل هذه العملية المركّبة.

  2. قد يكون هناك أكثر من عملية تنفيذ TensorFlow واحدة تستهدف مجموعة من الأدوات المدمَجة. عملية LiteRT. على سبيل المثال، هناك العديد من تنفيذات LSTM في TensorFlow (Keras، وBabelfish/lingvo، وما إلى ذلك) وتتألف كل منها من إلا أنه لا يزال من الممكن تحويلها جميعًا إلى عملية LSTM المدمجة نفسها في LiteRT.

وعلى هذا النحو، كان تحويل العمليات المدمجة هو تحديًا كبيرًا.

التفاف العملية المركّبة في 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 الشريحةer

تنفيذ العملية المدمجة كعملية 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) الاطلاع على مثال للدمج بشأن التضمين البحث. من الناحية النظرية، تحل كود التحويل محل العنصر المركب تنفيذ هذه الواجهة بالواجهة المدمجة.

في تمرير دوال التجهيز المركّبة، المكون الإضافي في تحويل الرمز.

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

التحويل إلى 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. تجميد جميع المتغيّرات للقراءة فقط

وهذا يتيح لنا تنفيذ عملية الدمج باستخدام الدوال التي تمثل العمليات المركبة قبل أن تؤدي وظائفها المضمنة والتجميد المتغير.

تنفيذ عملية دمج العمليات

لنلقِ نظرةً أكثر تفصيلاً على عملية fusion Pass. تفعل هذه البطاقة التالي:

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