Fusi operasi TensorFlow

Ringkasan

Halaman ini menjelaskan desain dan langkah-langkah yang diperlukan untuk mengonversi operasi komposit di TensorFlow menjadi operasi gabungan di LiteRT. Infrastruktur ini bersifat umum dan mendukung konversi operasi komposit apa pun di TensorFlow menjadi operasi gabungan yang sesuai di LiteRT.

Contoh penggunaan infrastruktur ini adalah fusi operasi RNN TensorFlow ke LiteRT, seperti yang dijelaskan di sini.

Apa yang dimaksud dengan operasi gabungan

menggambar

Operasi TensorFlow dapat berupa operasi primitif, misalnya tf.add atau dapat disusun dari operasi primitif lainnya, misalnya tf.einsum. Operasi primitif muncul sebagai satu node dalam grafik TensorFlow, sedangkan operasi komposit adalah kumpulan node dalam grafik TensorFlow. Mengeksekusi operasi komposit setara dengan mengeksekusi setiap operasi primitif konstituennya.

Operasi gabungan sesuai dengan satu operasi yang mencakup semua komputasi yang dilakukan oleh setiap operasi primitif dalam operasi komposit yang sesuai.

Manfaat operasi gabungan

Operasi gabungan ada untuk memaksimalkan performa penerapan kernel yang mendasarinya, dengan mengoptimalkan komputasi secara keseluruhan dan mengurangi jejak memori. Hal ini sangat berharga, terutama untuk beban kerja inferensi latensi rendah dan platform seluler dengan sumber daya terbatas.

Operasi gabungan juga menyediakan antarmuka tingkat yang lebih tinggi untuk menentukan transformasi kompleks seperti kuantisasi, yang jika tidak, tidak mungkin atau sangat sulit dilakukan pada tingkat yang lebih terperinci.

LiteRT memiliki banyak instance operasi gabungan karena alasan yang dijelaskan di atas. Operasi gabungan ini biasanya sesuai dengan operasi komposit dalam program TensorFlow sumber. Contoh operasi komposit di TensorFlow yang diimplementasikan sebagai satu operasi gabungan di LiteRT mencakup berbagai operasi RNN seperti LSTM urutan Unidirectional dan Bidirectional, konvolusi (conv2d, penambahan bias, relu), terhubung sepenuhnya (matmul, penambahan bias, relu), dan lainnya. Di LiteRT, kuantisasi LSTM saat ini hanya diimplementasikan dalam operasi LSTM gabungan.

Tantangan dengan operasi gabungan

Mengonversi operasi komposit dari TensorFlow ke operasi gabungan di LiteRT adalah masalah yang sulit. Hal ini dikarenakan:

  1. Operasi komposit ditampilkan dalam grafik TensorFlow sebagai serangkaian operasi primitif tanpa batas yang ditentukan dengan baik. Mengidentifikasi subgrafik yang sesuai dengan operasi komposit tersebut (misalnya, melalui pencocokan pola) bisa sangat sulit.

  2. Mungkin ada lebih dari satu penerapan TensorFlow yang menargetkan operasi fused LiteRT. Misalnya, ada banyak penerapan LSTM di TensorFlow (Keras, Babelfish/lingvo, dll.) dan masing-masing terdiri dari operasi primitif yang berbeda, tetapi semuanya masih dapat dikonversi ke operasi LSTM gabungan yang sama di LiteRT.

Oleh karena itu, konversi operasi gabungan terbukti cukup menantang.

Gabungkan operasi komposit dalam tf.function

Dalam banyak kasus, beberapa bagian model dapat dipetakan ke satu operasi di TFLite. Hal ini dapat membantu performa saat menulis penerapan yang dioptimalkan untuk operasi tertentu. Untuk dapat membuat operasi gabungan di TFLite, identifikasi bagian grafik yang merepresentasikan operasi gabungan dan bungkus dalam tf.function dengan atribut "experimental_implements" ke tf.function, yang memiliki nilai atribut tfl_fusable_op dengan nilai true. Jika operasi kustom memerlukan atribut, teruskan atribut tersebut sebagai bagian dari "experimental_implements" yang sama.

Contoh,

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))

Perhatikan bahwa Anda tidak perlu menetapkan allow_custom_ops pada konverter karena atribut tfl_fusable_op sudah menyiratkan hal ini.

Menerapkan operasi kustom dan mendaftar dengan Interpreter TFLite

Terapkan operasi gabungan Anda sebagai operasi Kustom TFLite - lihat petunjuk.

Perhatikan bahwa nama untuk mendaftarkan op harus mirip dengan nama yang ditentukan dalam atribut name dalam tanda tangan penerapan.

Contoh untuk op dalam contoh adalah

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

Mengonversi dari operasi komposit ke operasi gabungan (Advanced)

Arsitektur keseluruhan untuk mengonversi operasi komposit TensorFlow menjadi operasi gabungan LiteRT adalah sebagai berikut:

menggambar

Gabungkan operasi komposit dalam tf.function

Dalam kode sumber model TensorFlow, identifikasi dan abstrakkan operasi komposit ke dalam tf.function dengan anotasi fungsi experimental_implements. Lihat contoh pencarian penyematan. Fungsi menentukan antarmuka dan argumennya harus digunakan untuk menerapkan logika konversi.

Menulis kode konversi

Kode konversi ditulis per antarmuka fungsi dengan anotasi implements. Lihat contoh penggabungan untuk pencarian penyematan. Secara konseptual, kode konversi menggantikan penerapan gabungan antarmuka ini dengan penerapan gabungan.

Pada tahap prepare-composite-functions, masukkan kode konversi Anda.

Dalam penggunaan yang lebih canggih, Anda dapat menerapkan transformasi kompleks pada operan operasi komposit untuk mendapatkan operan operasi gabungan. Lihat kode konversi Keras LSTM sebagai contoh.

Mengonversi ke LiteRT

Gunakan API TFLiteConverter.from_saved_model untuk mengonversi ke LiteRT.

Di balik layar

Sekarang kita akan menjelaskan detail tingkat tinggi dari desain keseluruhan dalam mengonversi ke operasi gabungan di LiteRT.

Menyusun operasi di TensorFlow

Penggunaan tf.function dengan atribut fungsi experimental_implements memungkinkan pengguna menyusun operasi baru secara eksplisit menggunakan operasi primitif TensorFlow dan menentukan antarmuka yang diterapkan oleh operasi komposit yang dihasilkan. Hal ini sangat berguna karena memberikan:

  1. Batas yang ditentukan dengan baik untuk operasi komposit dalam grafik TensorFlow yang mendasarinya.
  2. Tentukan secara eksplisit antarmuka yang diimplementasikan oleh operasi ini. Argumen tf.function sesuai dengan argumen antarmuka ini.

Sebagai contoh, mari kita pertimbangkan operasi komposit yang ditentukan untuk menerapkan pencarian penyematan. Ini dipetakan ke operasi gabungan di 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

Dengan membuat model menggunakan operasi komposit melalui tf.function seperti yang diilustrasikan di atas, Anda dapat membangun infrastruktur umum untuk mengidentifikasi dan mengonversi operasi tersebut menjadi operasi LiteRT gabungan.

Memperluas konverter LiteRT

Konverter LiteRT yang dirilis awal tahun ini hanya mendukung pengimporan model TensorFlow sebagai grafik dengan semua variabel diganti dengan nilai konstanta yang sesuai. Hal ini tidak berfungsi untuk penggabungan operasi karena grafik tersebut memiliki semua fungsi inline sehingga variabel dapat diubah menjadi konstanta.

Untuk memanfaatkan tf.function dengan fitur experimental_implements selama proses konversi, fungsi harus dipertahankan hingga nanti dalam proses konversi.

Oleh karena itu, kami menerapkan alur kerja baru untuk mengimpor dan mengonversi model TensorFlow di pengonversi guna mendukung kasus penggunaan penggabungan operasi komposit. Secara khusus, fitur baru yang ditambahkan adalah:

  1. Mengimpor model tersimpan TensorFlow ke MLIR
  2. menggabungkan operasi komposit
  3. analisis perubahan variabel
  4. membekukan semua variabel hanya baca

Hal ini memungkinkan kita melakukan penggabungan operasi menggunakan fungsi yang merepresentasikan operasi komposit sebelum penyisipan fungsi dan pembekuan variabel.

Menerapkan penggabungan operasi

Mari kita lihat lebih detail proses penggabungan operasi. Kartu ini melakukan hal berikut:

  1. Lakukan loop melalui semua fungsi dalam modul MLIR.
  2. Jika suatu fungsi memiliki atribut tf._implements, berdasarkan nilai atribut, panggil utilitas penggabungan operasi yang sesuai.
  3. Utilitas penggabungan operasi beroperasi pada operand dan atribut fungsi (yang berfungsi sebagai antarmuka untuk konversi) dan mengganti isi fungsi dengan isi fungsi yang setara yang berisi operasi gabungan.
  4. Dalam banyak kasus, isi yang diganti akan berisi operasi selain operasi gabungan. Transformasi ini sesuai dengan beberapa transformasi statis pada operand fungsi untuk mendapatkan operand operasi gabungan. Karena semua komputasi ini dapat dilipat secara konstan, komputasi ini tidak akan ada di flatbuffer yang diekspor, yang hanya akan berisi operasi gabungan.

Berikut adalah cuplikan kode dari kartu yang menunjukkan alur kerja utama:

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 */
}

Berikut cuplikan kode yang menunjukkan pemetaan operasi komposit ini ke operasi gabungan di LiteRT yang memanfaatkan fungsi sebagai antarmuka konversi.

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