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, बायस ऐड, relu), पूरी तरह से कनेक्टेड (matmul, बायस ऐड, relu) और बहुत कुछ . TensorFlow Lite में, LSTM परिमाणीकरण वर्तमान में केवल फ़्यूज्ड LSTM संचालन में लागू किया गया है।

जुड़े हुए संचालन के साथ चुनौतियाँ

TensorFlow Lite में समग्र संचालन को TensorFlow से फ़्यूज्ड संचालन में परिवर्तित करना एक कठिन समस्या है। यह है क्योंकि:

  1. कंपोजिट ऑपरेशंस को TensorFlow ग्राफ़ में एक अच्छी तरह से परिभाषित सीमा के बिना आदिम ऑपरेशंस के एक सेट के रूप में दर्शाया गया है। ऐसे समग्र ऑपरेशन के अनुरूप उप-ग्राफ़ की पहचान करना (उदाहरण के लिए पैटर्न मिलान के माध्यम से) बहुत चुनौतीपूर्ण हो सकता है।

  2. फ़्यूज्ड TensorFlow Lite ऑपरेशन को लक्षित करने वाले एक से अधिक TensorFlow कार्यान्वयन हो सकते हैं। उदाहरण के लिए, TensorFlow (Keras, Babelfish/lingvo आदि) में कई LSTM कार्यान्वयन हैं और इनमें से प्रत्येक अलग-अलग आदिम संचालन से बना है, लेकिन उन सभी को अभी भी TensorFlow Lite में एक ही फ़्यूज्ड LSTM ऑपरेशन में परिवर्तित किया जा सकता है।

इस प्रकार, फ़्यूज़्ड ऑपरेशनों का रूपांतरण काफी चुनौतीपूर्ण साबित हुआ है।

कंपोजिट ऑपरेशन को tf.function में लपेटें

कई मामलों में, मॉडल के कुछ हिस्से को TFLite में एकल ऑपरेशन में मैप किया जा सकता है। विशिष्ट परिचालनों के लिए अनुकूलित कार्यान्वयन लिखते समय यह प्रदर्शन में मदद कर सकता है। TFLite में फ़्यूज़्ड ऑपरेशन बनाने में सक्षम होने के लिए, ग्राफ़ के उस हिस्से की पहचान करें जो फ़्यूज़्ड ऑपरेशन का प्रतिनिधित्व करता है और इसे 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 दुभाषिया के साथ पंजीकरण करें

अपने फ़्यूज्ड ऑपरेशन को 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 कंपोजिट ऑपरेशंस को TensorFlow Lite फ़्यूज्ड ऑपरेशंस में परिवर्तित करने के लिए समग्र आर्किटेक्चर नीचे है:

चित्रकला

कंपोजिट ऑपरेशन को tf.function में लपेटें

TensorFlow मॉडल स्रोत कोड में, प्रयोगात्मक_इम्प्लीमेंट्स फ़ंक्शन एनोटेशन के साथ समग्र ऑपरेशन को tf.function में पहचानें और सारबद्ध करें। एम्बेडिंग लुकअप का एक उदाहरण देखें. फ़ंक्शन इंटरफ़ेस को परिभाषित करता है और रूपांतरण तर्क को लागू करने के लिए इसके तर्कों का उपयोग किया जाना चाहिए।

रूपांतरण कोड लिखें

रूपांतरण कोड implements एनोटेशन के साथ फ़ंक्शन के इंटरफ़ेस के अनुसार लिखा जाता है। लुकअप को एम्बेड करने के लिए फ़्यूज़न का एक उदाहरण देखें। वैचारिक रूप से, रूपांतरण कोड इस इंटरफ़ेस के समग्र कार्यान्वयन को फ़्यूज्ड इंटरफ़ेस से बदल देता है।

प्रिपेयर-कंपोजिट-फंक्शन पास में, अपने रूपांतरण कोड में प्लगइन करें।

अधिक उन्नत उपयोगों में, फ़्यूज्ड ऑपरेशन के ऑपरेंड को प्राप्त करने के लिए समग्र ऑपरेशन के ऑपरेंड के जटिल परिवर्तनों को लागू करना संभव है। केरस एलएसटीएम देखें। एक उदाहरण के रूप में रूपांतरण कोड.

TensorFlow लाइट में कनवर्ट करें

TensorFlow Lite में कनवर्ट करने के लिए TFLiteConverter.from_saven_model API का उपयोग करें।

हुड के नीचे

अब हम TensorFlow Lite में फ़्यूज्ड संचालन में परिवर्तित करने में समग्र डिज़ाइन के उच्च स्तरीय विवरण का वर्णन करते हैं।

TensorFlow में कंपोज़िंग ऑपरेशन

प्रयोगात्मक_इम्प्लीमेंट्स फ़ंक्शन विशेषता के साथ tf.function का उपयोग उपयोगकर्ताओं को 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 मॉडल को ग्राफ़ के रूप में आयात करने का समर्थन करता है, जिसमें सभी चर को उनके संबंधित स्थिर मानों के साथ बदल दिया जाता है। यह ऑपरेशन फ़्यूज़न के लिए काम नहीं करता है क्योंकि ऐसे ग्राफ़ में सभी फ़ंक्शन इनलाइन होते हैं ताकि चर को स्थिरांक में बदला जा सके।

रूपांतरण प्रक्रिया के दौरान experimental_implements सुविधा के साथ tf.function लाभ उठाने के लिए, फ़ंक्शन को रूपांतरण प्रक्रिया में बाद तक संरक्षित रखने की आवश्यकता होती है।

इस प्रकार, हमने समग्र ऑपरेशन फ़्यूज़न उपयोग के मामले का समर्थन करने के लिए कनवर्टर में टेन्सरफ्लो मॉडल को आयात और परिवर्तित करने का एक नया वर्कफ़्लो लागू किया। विशेष रूप से, जोड़ी गई नई सुविधाएँ हैं:

  1. TensorFlow सहेजे गए मॉडल को MLIR में आयात करना
  2. फ़्यूज़ समग्र संचालन
  3. परिवर्तनीय परिवर्तनशीलता विश्लेषण
  4. सभी रीड-ओनली वेरिएबल्स को फ़्रीज़ करें

यह हमें फ़ंक्शन इनलाइनिंग और वेरिएबल फ्रीजिंग से पहले समग्र संचालन का प्रतिनिधित्व करने वाले कार्यों का उपयोग करके ऑपरेशन फ़्यूज़न करने की अनुमति देता है।

ऑपरेशन फ़्यूज़न का कार्यान्वयन

आइए ऑपरेशन फ़्यूज़न पास को अधिक विस्तार से देखें। यह पास निम्नलिखित कार्य करता है:

  1. एमएलआईआर मॉड्यूल में सभी कार्यों के माध्यम से लूप करें।
  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 */
}

यहां कोड स्निपेट है जो इस समग्र ऑपरेशन को रूपांतरण इंटरफ़ेस के रूप में फ़ंक्शन का लाभ उठाते हुए 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());
  }