همجوشی عملیات تنسورفلو

نمای کلی

این صفحه طراحی و مراحل مورد نیاز برای تبدیل عملیات ترکیبی در TensorFlow به عملیات ذوب شده در LiteRT را شرح می‌دهد. این زیرساخت یک هدف کلی است و از تبدیل هر عملیات ترکیبی در TensorFlow به یک عملیات ذوب شده مربوطه در LiteRT پشتیبانی می کند.

نمونه ای از استفاده از این زیرساخت، ادغام عملیات RNN TensorFlow به LiteRT است که در اینجا به تفصیل شرح داده شده است.

عملیات ذوب شده چیست؟

عملیات TensorFlow یا می تواند ops اولیه باشد به عنوان مثال tf.add یا می تواند از سایر عملیات های اولیه مانند tf.einsum تشکیل شود. یک عملیات اولیه به صورت یک گره در گراف TensorFlow نشان داده می شود در حالی که یک عملیات ترکیبی مجموعه ای از گره ها در گراف TensorFlow است. اجرای یک عملیات ترکیبی معادل اجرای هر یک از عملیات اولیه تشکیل دهنده آن است.

یک عملیات ذوب شده مربوط به یک عملیات واحد است که تمام محاسبات انجام شده توسط هر عملیات اولیه را در عملیات ترکیبی مربوطه جمع می کند.

مزایای عملیات ذوب شده

عملیات ذوب شده برای به حداکثر رساندن عملکرد پیاده سازی هسته اصلی خود، با بهینه سازی محاسبات کلی و کاهش ردپای حافظه وجود دارد. این بسیار ارزشمند است، به ویژه برای بارهای کاری استنتاج با تأخیر کم و پلتفرم های تلفن همراه با محدودیت منابع.

عملیات ذوب شده همچنین یک رابط سطح بالاتر برای تعریف تبدیل های پیچیده مانند کوانتیزاسیون ارائه می دهد که در غیر این صورت انجام آن در یک سطح دانه ای تر غیر ممکن یا بسیار دشوار است.

LiteRT به دلایلی که در بالا توضیح داده شد، موارد زیادی از عملیات ذوب شده دارد. این عملیات ذوب شده معمولاً با عملیات ترکیبی در برنامه TensorFlow منبع مطابقت دارد. نمونه‌هایی از عملیات ترکیبی در TensorFlow که به صورت یک عملیات ذوب شده در LiteRT پیاده‌سازی می‌شوند شامل عملیات‌های مختلف RNN مانند دنباله تک‌جهته و دو جهته LSTM، کانولوشن (conv2d، bias add، relu)، کاملا متصل (matmul، bias add، 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 ثبت نام کنید

عملیات ذوب شده خود را به عنوان یک عملیات سفارشی 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، عملیات ترکیبی را با حاشیه نویسی تابع Experimental_implements به یک tf.function انتزاع و انتزاع کنید. نمونه ای از جستجوی جاسازی را ببینید. تابع رابط را تعریف می کند و آرگومان های آن باید برای پیاده سازی منطق تبدیل استفاده شوند.

کد تبدیل را بنویسید

کد تبدیل در رابط تابع با حاشیه نویسی implements نوشته می شود. یک نمونه ترکیبی برای جستجوی جاسازی را ببینید. از نظر مفهومی، کد تبدیل، اجرای ترکیبی این رابط را با رابط فیوژن جایگزین می‌کند.

در گذرنامه آماده-کامپوزیت-توابع، کد تبدیل خود را افزونه کنید.

در کاربردهای پیشرفته تر، می توان تبدیل های پیچیده عملوندهای عملیات ترکیبی را به منظور استخراج عملوندهای عملیات ذوب شده پیاده سازی کرد. Keras LSTM را ببینید. کد تبدیل به عنوان مثال

تبدیل به LiteRT

برای تبدیل به LiteRT از TFLiteConverter.from_saved_model API استفاده کنید.

زیر کاپوت

اکنون جزئیات سطح بالایی از طراحی کلی را در تبدیل به عملیات ذوب شده در LiteRT شرح می دهیم.

نوشتن عملیات در TensorFlow

استفاده از tf.function با ویژگی تابع Experimental_implements به کاربران این امکان را می دهد که به صراحت عملیات جدید را با استفاده از عملیات اولیه TensorFlow ترکیب کرده و رابطی را که عملیات ترکیبی حاصل اجرا می کند، مشخص کنند. این بسیار مفید است زیرا ارائه می دهد:

1. یک مرز کاملاً تعریف شده برای عملیات ترکیبی در گراف TensorFlow زیرین.
2. به صراحت رابطی را که این عملیات پیاده سازی می کند مشخص کنید. آرگومان های تابع tf. با آرگومان های این رابط مطابقت دارد.

به عنوان مثال، اجازه دهید یک عملیات ترکیبی تعریف شده برای پیاده سازی جستجوی جاسازی را در نظر بگیریم. این به یک عملیات ذوب شده در 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 را داشته باشد، بر اساس مقدار مشخصه، عملیات مناسب fusion utility را فراخوانی می کند.
3. ابزار عملیات 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 */
}

در اینجا قطعه کدی است که نگاشت این عملیات ترکیبی را به یک عملیات ذوب شده در 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());
  }