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

نمای کلی

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

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

عملیات ترکیبی چیست؟

عملیات TensorFlow می‌توانند عملیات اولیه باشند، مثلاً 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 در امضای implements مشخص شده است.

مثالی برای op در مثال این است

  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 نوشته می‌شود. به مثالی از ترکیب برای embedding lookup مراجعه کنید. از نظر مفهومی، کد تبدیل، پیاده‌سازی ترکیبی این رابط را با پیاده‌سازی ترکیب‌شده جایگزین می‌کند.

در پاسِ prepare-composite-functions، کد تبدیل خود را اضافه کنید.

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

تبدیل به LiteRT

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

زیر کاپوت

اکنون جزئیات سطح بالای طراحی کلی را در تبدیل به عملیات ترکیبی در 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. تمام متغیرهای فقط خواندنی را مسدود می‌کند

این به ما اجازه می‌دهد تا قبل از inline کردن تابع و ثابت کردن متغیر، ادغام عملیات را با استفاده از توابعی که نشان‌دهنده عملیات ترکیبی هستند، انجام دهیم.

پیاده‌سازی ادغام عملیات

بیایید با جزئیات بیشتری به مرحله‌ی فیوژن عملیات نگاه کنیم. این مرحله موارد زیر را انجام می‌دهد:

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());
  }