TensorFlow ऑपरेशन फ़्यूज़न

खास जानकारी

इस पेज पर, TensorFlow की कंपोज़िट ऑपरेशन को TensorFlow Lite के फ़्यूज़ किए गए ऑपरेशन में बदलने के लिए ज़रूरी डिज़ाइन और चरणों के बारे में बताया गया है. इस इन्फ़्रास्ट्रक्चर का इस्तेमाल आम तौर पर किया जाता है और यह TensorFlow में किसी भी कंपोज़िट ऑपरेशन को TensorFlow Lite में संबंधित फ़्यूज़ ऑपरेशन में बदलने की सुविधा देता है.

इस इंफ़्रास्ट्रक्चर का एक उदाहरण है, TensorFlow RNN ऑपरेशन फ़्यूज़न, TensorFlow Lite. इसके बारे में यहां बताया गया है.

फ़्यूज़ की गई संक्रियाएं क्या हैं

ड्रॉइंग

TensorFlow की कार्रवाइयां या तो शुरुआती कार्रवाई हो सकती हैं, जैसे कि tf.add या फिर tf.einsum जैसी दूसरी प्रोसेस से, इन्हें कॉन्फ़िगर किया जा सकता है. एक प्रिमिटिव ऑपरेशन, TensorFlow ग्राफ़ में एक नोड के तौर पर दिखता है, जबकि एक कंपोज़िट ऑपरेशन, TensorFlow ग्राफ़ में नोड का कलेक्शन होता है. कंपोज़िट ऑपरेशन को एक्ज़ीक्यूट करना, उसके हर शुरुआती ऑपरेशन को लागू करने के बराबर है.

कई चीज़ों के आधार पर की गई कार्रवाई, एक ऐसी कार्रवाई होती है जो हर प्रिमिटिव ऑपरेशन से की जाने वाली सभी कंप्यूटिंग को शामिल करती है.

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

कई तरह की कार्रवाइयों का इस्तेमाल किया जाता है, ताकि पूरे कंप्यूटेशन को ऑप्टिमाइज़ करके और मेमोरी फ़ुटप्रिंट को कम करके, उनके कर्नेल को लागू करने की प्रोसेस की परफ़ॉर्मेंस को बेहतर बनाया जा सके. खास तौर पर, इंतज़ार का समय कम करने वाले वर्कलोड और सीमित संसाधन वाले मोबाइल प्लैटफ़ॉर्म के लिए यह बहुत अहम है.

फ़्यूज़्ड ऑपरेशन, क्वांटाइज़ेशन जैसे जटिल ट्रांसफ़ॉर्मेशन को परिभाषित करने के लिए हाई लेवल का इंटरफ़ेस भी उपलब्ध कराते हैं. हालांकि, ऐसा करना मुमकिन नहीं होता या ज़्यादा जानकारी के लिए ऐसा करना बहुत मुश्किल होता.

TensorFlow Lite में ऊपर बताई गई वजहों से, फ़्यूज़ की गई कार्रवाइयां कई बार हुई हैं. आम तौर पर, ये फ़्यूज़ किए गए ऑपरेशन TensorFlow प्रोग्राम में होने वाले कंपोज़िट ऑपरेशन के हिसाब से होते हैं. TensorFlow में TensorFlow के कंपोज़िट ऑपरेशन के उदाहरण TensorFlow Lite में, फ़िलहाल एलएसटीएम क्वांटाइज़ेशन की सुविधा सिर्फ़ फ़्यूज़ की गई एलएसटीएम कार्रवाइयों में काम करती है.

कई चीज़ों के एक साथ इस्तेमाल होने से जुड़ी चुनौतियां

TensorFlow Lite में कंपोज़िट ऑपरेशन को TensorFlow से फ़्यूज़ ऑपरेशन में बदलना एक कठिन समस्या है. ऐसा इसलिए होता है, क्योंकि:

  1. TensorFlow ग्राफ़ में, कंपोज़िट ऑपरेशन को प्रिमिटिव ऑपरेशन के एक सेट के तौर पर दिखाया जाता है. हालांकि, इसकी सीमा तय नहीं होती है. ऐसे कंपोज़िट ऑपरेशन से जुड़े सब-ग्राफ़ को पहचानना बहुत मुश्किल हो सकता है (जैसे पैटर्न मैचिंग के ज़रिए).

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

इसलिए, आपस में जुड़े हुए ऑपरेशनों का रूपांतरण काफ़ी चुनौती भरा साबित हुआ है.

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

कई मामलों में, मॉडल के कुछ हिस्से को TFLite में एक ही कार्रवाई के लिए मैप किया जा सकता है. इससे किसी खास ऑपरेशन के लिए ऑप्टिमाइज़ किया गया तरीका लागू करते समय, परफ़ॉर्मेंस बेहतर करने में मदद मिल सकती है. TFLite में एक मिलाया गया ऑपरेशन बनाने के लिए, ग्राफ़ के उस हिस्से की पहचान करें जो फ़्यूज़ किए गए ऑपरेशन को दिखाता है और उसे tf.function में, ऐसे tf.function में रैप करें जिसमें "experimental_applieds" एट्रिब्यूट शामिल हैं, जिसमें true वैल्यू के साथ एट्रिब्यूट वैल्यू tfl_fusable_op है. अगर कस्टम ऑपरेशन में एट्रिब्यूट इस्तेमाल होते हैं, तो उन्हें उसी "experimental_applieds" के हिस्से के तौर पर पास करें.

उदाहरण,

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

कन्वर्ज़न कोड लिखें

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

तैयारी-कंपोज़िट-फ़ंक्शन पास में, अपने कन्वर्ज़न कोड का प्लग इन करें.

ज़्यादा बेहतर इस्तेमाल के मामले में, कंपोज़िट ऑपरेशन के ऑपरेंड के जटिल ट्रांसफ़ॉर्मेशन को लागू करना मुमकिन है, ताकि कई फ़्यूज़ किए गए ऑपरेशन के बारे में पता लगाया जा सके. Keras LSTM देखें. उदाहरण के लिए, कन्वर्ज़न कोड.

TensorFlow Lite में बदलें

TensorFlow Lite में बदलने के लिए, TFLiteConverter.from_saved_model एपीआई का इस्तेमाल करें.

हुड के नीचे

अब हम TensorFlow Lite में, फ़्यूज़ किए गए ऑपरेशन में बदलने के पूरे डिज़ाइन के बारे में ज़्यादा जानकारी देते हैं.

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

experimental_implements फ़ंक्शन के साथ 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 का फ़ायदा पाने के लिए, फ़ंक्शन को कन्वर्ज़न प्रोसेस में, बाद तक सुरक्षित रखा जाना चाहिए.

इसलिए, हमने कन्वर्टर में TensorFlow के मॉडल इंपोर्ट और बदलने के लिए एक नया वर्कफ़्लो लागू किया है, ताकि कंपोज़िट ऑपरेशन फ़्यूज़न के इस्तेमाल के उदाहरण के साथ काम किया जा सके. खास तौर पर, जोड़ी गई नई सुविधाएं ये हैं:

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

इससे हमें फ़ंक्शन इनलाइनिंग और वैरिएबल फ़्रीज़िंग से पहले कंपोज़िट ऑपरेशन दिखाने वाले फ़ंक्शन का इस्तेमाल करके, ऑपरेशन फ़्यूज़न करने में मदद मिलती है.

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

आइए, 'ऑपरेशन फ़्यूज़न पास' के बारे में ज़्यादा जानकारी देखते हैं. इस पास से ये काम किए जाते हैं:

  1. MLIR मॉड्यूल में सभी फ़ंक्शन को लूप में चलाएं.
  2. अगर किसी फ़ंक्शन में tf._sync एट्रिब्यूट है, तो एट्रिब्यूट की वैल्यू के आधार पर, सही ऑपरेशन फ़्यूज़न यूटिलिटी को कॉल करता है.
  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());
  }