Penggabungan operasi TensorFlow

Ringkasan

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

Contoh penggunaan infrastruktur ini adalah penggabungan operasi TensorFlow RNN ke TensorFlow Lite, seperti yang dirinci di sini .

Apa itu operasi gabungan

menggambar

Operasi TensorFlow dapat berupa operasi primitif misalnya tf.add atau dapat terdiri dari operasi primitif lainnya misalnya tf.einsum . Operasi primitif muncul sebagai satu node di grafik TensorFlow, sedangkan operasi gabungan adalah kumpulan node di grafik TensorFlow. Mengeksekusi operasi gabungan sama dengan mengeksekusi setiap operasi primitif penyusunnya.

Operasi gabungan berhubungan dengan operasi tunggal yang menggabungkan semua komputasi yang dilakukan oleh setiap operasi primitif dalam operasi gabungan terkait.

Manfaat operasi gabungan

Operasi gabungan ada untuk memaksimalkan kinerja implementasi kernel yang mendasarinya, dengan mengoptimalkan keseluruhan komputasi 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 mendefinisikan transformasi kompleks seperti kuantisasi, yang jika tidak dilakukan akan menjadi tidak mungkin atau sangat sulit dilakukan pada tingkat yang lebih terperinci.

TensorFlow Lite memiliki banyak contoh operasi gabungan karena alasan yang disebutkan di atas. Operasi gabungan ini biasanya berhubungan dengan operasi gabungan dalam program sumber TensorFlow. Contoh operasi komposit di TensorFlow yang diimplementasikan sebagai operasi gabungan tunggal di TensorFlow Lite mencakup berbagai operasi RNN seperti LSTM rangkaian Searah dan Dua Arah, konvolusi (konv2d, penambahan bias, relu), terhubung sepenuhnya (matmul, penambahan bias, relu) dan banyak lagi . Di TensorFlow Lite, kuantisasi LSTM saat ini hanya diterapkan dalam operasi LSTM yang digabungkan.

Tantangan dengan operasi yang menyatu

Mengonversi operasi gabungan dari TensorFlow menjadi operasi gabungan di TensorFlow Lite adalah masalah yang sulit. Hal ini dikarenakan:

  1. Operasi gabungan direpresentasikan dalam grafik TensorFlow sebagai serangkaian operasi primitif tanpa batas yang jelas. Akan sangat sulit untuk mengidentifikasi (misalnya melalui pencocokan pola) sub-grafik yang terkait dengan operasi gabungan tersebut.

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

Oleh karena itu, konversi operasi fusi terbukti cukup menantang.

Bungkus operasi komposit dalam tf.function

Dalam banyak kasus, beberapa bagian model dapat dipetakan ke satu operasi di TFLite. Hal ini dapat membantu kinerja saat menulis implementasi yang dioptimalkan untuk operasi tertentu. Untuk dapat membuat operasi gabungan di TFLite, identifikasi bagian grafik yang mewakili operasi gabungan dan gabungkan dalam tf.function dengan atribut "experimental_implements" ke tf.function , yang memiliki nilai atribut tfl_fusable_op dengan nilai true . Jika operasi khusus mengambil atribut, teruskan atribut tersebut sebagai bagian dari "implementasi_eksperimental" 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 allow_custom_ops pada konverter karena atribut tfl_fusable_op sudah menyiratkan hal ini.

Terapkan operasi khusus dan daftar dengan TFLite Interpreter

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

Perhatikan bahwa, nama untuk mendaftarkan operasi harus serupa dengan nama yang ditentukan dalam atribut name di tanda tangan implementasi.

Contoh op pada contoh tersebut 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 leburan (Lanjutan)

Arsitektur keseluruhan untuk mengonversi operasi gabungan TensorFlow ke operasi gabungan TensorFlow Lite adalah sebagai berikut:

menggambar

Bungkus operasi komposit dalam tf.function

Dalam kode sumber model TensorFlow, identifikasi dan abstraksi operasi komposit menjadi tf.function dengan anotasi fungsi eksperimental_implements . Lihat contoh penyematan pencarian . Fungsi tersebut mendefinisikan antarmuka dan argumennya harus digunakan untuk mengimplementasikan logika konversi.

Tulis kode konversi

Kode konversi ditulis per antarmuka fungsi dengan anotasi implements . Lihat contoh fusi untuk menyematkan pencarian . Secara konseptual, kode konversi menggantikan implementasi gabungan antarmuka ini dengan implementasi gabungan.

Di pass fungsi persiapan-komposit, masukkan kode konversi Anda.

Dalam penggunaan yang lebih maju, dimungkinkan untuk menerapkan transformasi kompleks dari operan operasi gabungan untuk mendapatkan operan dari operasi gabungan. Lihat Keras LSTM . kode konversi sebagai contoh.

Konversikan ke TensorFlow Lite

Gunakan TFLiteConverter.from_saved_model API untuk mengonversi ke TensorFlow Lite.

Dibawah tenda

Kami sekarang menjelaskan detail tingkat tinggi dari keseluruhan desain dalam mengonversi operasi gabungan di TensorFlow Lite.

Menulis operasi di TensorFlow

Penggunaan tf.function dengan atribut fungsi eksperimental_implements memungkinkan pengguna membuat operasi baru secara eksplisit menggunakan operasi primitif TensorFlow dan menentukan antarmuka yang diimplementasikan oleh operasi gabungan yang dihasilkan. Ini sangat berguna karena menyediakan:

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

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

Dengan membuat model menggunakan operasi komposit melalui tf.function seperti yang diilustrasikan di atas, infrastruktur umum dapat dibangun untuk mengidentifikasi dan mengonversi operasi tersebut menjadi operasi TensorFlow Lite yang digabungkan.

Memperluas konverter TensorFlow Lite

Konverter TensorFlow Lite yang dirilis awal tahun ini hanya mendukung impor model TensorFlow sebagai grafik dengan semua variabel diganti dengan nilai konstanta yang sesuai. Ini tidak berlaku untuk operasi fusi karena grafik tersebut memiliki semua fungsi yang sejajar sehingga variabel dapat diubah menjadi konstanta.

Untuk memanfaatkan tf.function dengan fitur experimental_implements selama proses konversi, fungsi-fungsi tersebut perlu dipertahankan hingga proses konversi nanti.

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

  1. Mengimpor model tersimpan TensorFlow ke MLIR
  2. sekering operasi komposit
  3. analisis mutabilitas variabel
  4. bekukan semua variabel read-only

Hal ini memungkinkan kita untuk melakukan fusi operasi menggunakan fungsi yang mewakili operasi komposit sebelum fungsi inlining dan pembekuan variabel.

Menerapkan fusi operasi

Mari kita lihat operasi fusion pass lebih detail. Pass ini melakukan hal berikut:

  1. Ulangi semua fungsi dalam modul MLIR.
  2. Jika suatu fungsi memiliki atribut tf._implements, berdasarkan nilai atribut, panggil utilitas fusi operasi yang sesuai.
  3. Utilitas fusi operasi beroperasi pada operan dan atribut fungsi (yang berfungsi sebagai antarmuka untuk konversi) dan menggantikan isi fungsi dengan isi fungsi setara yang berisi operasi fusi.
  4. Dalam banyak kasus, bodi yang diganti akan berisi operasi selain operasi fusi. Ini sesuai dengan beberapa transformasi statis pada operan fungsi untuk mendapatkan operan dari operasi gabungan. Karena semua penghitungan ini dapat dilipat secara konstan, penghitungan tersebut tidak akan ada di flatbuffer yang diekspor di mana hanya operasi gabungan yang ada.

Berikut cuplikan kode dari pass 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 gabungan ini ke operasi gabungan di TensorFlow Lite 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());
  }