Bantuan melindungi Great Barrier Reef dengan TensorFlow pada Kaggle Bergabung Tantangan

Fusi operasi TensorFlow

Gambaran

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

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

Apa itu operasi fusi

gambar

Operasi TensorFlow dapat berupa operasi primitif mis. Tf.add atau dapat disusun dari operasi primitif lain mis. Tf.einsum . Operasi primitif muncul sebagai node tunggal dalam grafik TensorFlow, sedangkan operasi gabungan adalah kumpulan node dalam grafik TensorFlow. Menjalankan operasi gabungan sama dengan menjalankan setiap operasi primitif penyusunnya.

Operasi gabungan sesuai dengan operasi tunggal yang memasukkan semua komputasi yang dilakukan oleh setiap operasi primitif dalam operasi gabungan yang sesuai.

Manfaat operasi fusi

Operasi gabungan ada untuk memaksimalkan kinerja implementasi kernel yang mendasarinya, dengan mengoptimalkan komputasi keseluruhan dan mengurangi jejak memori. Ini sangat berharga, terutama untuk beban kerja inferensi latensi rendah dan platform seluler dengan sumber daya yang terbatas.

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

TensorFlow Lite memiliki banyak contoh operasi gabungan karena alasan yang diartikulasikan di atas. Operasi fusi ini biasanya sesuai dengan operasi gabungan dalam program TensorFlow sumber. Contoh operasi gabungan di TensorFlow yang diimplementasikan sebagai operasi gabungan tunggal di TensorFlow Lite mencakup berbagai operasi RNN seperti LSTM urutan 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 gabungan.

Tantangan dengan operasi gabungan

Mengonversi operasi komposit dari TensorFlow menjadi operasi fusi di TensorFlow Lite adalah masalah yang sulit. Hal ini karena:

  1. Operasi komposit direpresentasikan dalam grafik TensorFlow sebagai rangkaian operasi primitif tanpa batas yang ditentukan dengan baik. Ini bisa sangat menantang untuk mengidentifikasi (misalnya melalui pencocokan pola) sub-grafik yang sesuai 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 terdiri dari operasi primitif yang berbeda, tetapi semuanya masih dapat dikonversi ke operasi LSTM gabungan yang sama di TensorFlow Lite.

Dengan demikian, konversi operasi fusi terbukti cukup menantang.

Mengonversi dari operasi gabungan ke operasi gabungan

Arsitektur keseluruhan untuk mengonversi operasi komposit TensorFlow menjadi operasi fusi TensorFlow Lite adalah di bawah ini:

gambar

Bungkus operasi komposit dalam fungsi tf.function

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

Tulis kode konversi

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

Di pass persiapan-fungsi-komposit, plugin di kode konversi Anda.

Dalam penggunaan yang lebih maju, dimungkinkan untuk mengimplementasikan transformasi kompleks dari operan operasi komposit untuk mendapatkan operand dari operasi fusi. Lihat Keras LSTM . kode konversi sebagai contoh.

Ubah ke TensorFlow Lite

Gunakan TFLiteConverter.from_saved_model API untuk mengonversi ke TensorFlow Lite.

Dibawah tenda

Sekarang kami mendeskripsikan detail tingkat tinggi dari keseluruhan desain dalam mengonversi menjadi operasi gabungan di TensorFlow Lite.

Operasi penulisan di TensorFlow

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

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

Sebagai contoh, mari pertimbangkan operasi gabungan yang ditentukan untuk mengimplementasikan pencarian embedding. Ini memetakan 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 dibuat 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 pengimporan model TensorFlow sebagai grafik dengan semua variabel diganti dengan nilai konstanta yang sesuai. Ini tidak bekerja untuk fusi operasi karena grafik seperti itu memiliki semua fungsi sebaris sehingga variabel dapat diubah menjadi konstanta.

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

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

  1. Mengimpor model tersimpan TensorFlow ke MLIR
  2. operasi komposit sekering
  3. analisis mutabilitas variabel
  4. membekukan semua variabel hanya-baca

Ini memungkinkan kita untuk melakukan fusi operasi menggunakan fungsi yang mewakili operasi komposit sebelum fungsi sebaris dan pembekuan variabel.

Menerapkan fusi operasi

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

  1. Ulangi semua fungsi dalam modul MLIR.
  2. Jika suatu fungsi memiliki atribut tf._implements, berdasarkan nilai atributnya, memanggil utilitas fusi operasi yang sesuai.
  3. Utilitas fusi operasi beroperasi pada operand dan atribut fungsi (yang berfungsi sebagai antarmuka untuk konversi) dan menggantikan badan fungsi dengan badan fungsi ekivalen yang berisi operasi gabungan.
  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 fusi. Karena semua penghitungan ini dapat secara konstan dilipat, mereka tidak akan ada di flatbuffer yang diekspor di mana hanya operasi fusi yang akan ada.

Berikut adalah potongan 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 komposit 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());
  }