Konversi RNN TensorFlow ke TensorFlow Lite

Ringkasan

TensorFlow Lite mendukung konversi model RNN TensorFlow ke operasi LSTM gabungan TensorFlow Lite. Operasi gabungan diterapkan untuk memaksimalkan performa implementasi kernel yang mendasarinya, serta memberikan antarmuka dengan level yang lebih tinggi untuk menentukan transformasi kompleks seperti kuantisasi.

Karena ada banyak varian RNN API di TensorFlow, pendekatan kami menjadi dua kali lipat:

  1. Berikan dukungan native untuk TensorFlow RNN API standar seperti Keras LSTM. Opsi ini direkomendasikan.
  2. Menyediakan antarmuka ke dalam infrastruktur konversi untuk implementasi RNN ditentukan oleh pengguna untuk dihubungkan dan dikonversi ke TensorFlow Lite. Kami menyediakan beberapa contoh langsung dari konversi tersebut menggunakan antarmuka RNN LSTMCellSimple dan LayerNormalizedLSTMCellSimple lingvo.

API Konverter

Fitur ini merupakan bagian dari rilis TensorFlow 2.3. Kode ini juga tersedia melalui pip tf-nightly atau dari head.

Fungsionalitas konversi ini tersedia saat melakukan konversi ke TensorFlow Lite melalui SaveModel atau langsung dari model Keras. Lihat contoh penggunaan.

Dari model tersimpan

# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)

# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

Dari model Keras

# build a Keras model
keras_model = build_keras_lstm(...)

# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

Contoh

Keras LSTM ke TensorFlow Lite Colab mengilustrasikan penggunaan menyeluruh dengan penafsir TensorFlow Lite.

TensorFlow RNN API didukung

Kami mendukung konversi siap pakai dari Keras LSTM ke TensorFlow Lite. Untuk mengetahui detail tentang cara kerjanya, lihat antarmuka Keras LSTM dan logika konversi di sini.

Penting juga untuk menyoroti kontrak LSTM TensorFlow Lite sehubungan dengan definisi operasi Keras:

  1. Dimensi 0 dari tensor input adalah ukuran batch.
  2. Dimensi 0 dari tensor recurrent_weight adalah jumlah output.
  3. Tensor weight dan recurrent_kernel ditransposisikan.
  4. Tensor bobot yang ditransposisikan, recurrent_kernel yang ditransposisikan, dan tensor bias dibagi menjadi 4 tensor dengan ukuran yang sama di sepanjang dimensi 0. Ini berhubungan dengan gate input, remember gate, sel, dan gate output.

Varian Keras LSTM

Jurusan waktu

Pengguna dapat memilih besar waktu atau tidak ada waktu besar. Keras LSTM menambahkan atribut utama waktu pada atribut def fungsi. Untuk LSTM urutan searah, kita dapat memetakan ke atribut utama waktu unidirecional_sequence_lstm.

LSTM Dua Arah

LSTM dua arah dapat diimplementasikan dengan dua lapisan Keras LSTM, satu untuk maju dan satu lagi untuk mundur. Lihat contohnya di sini. Setelah melihat atribut go_backward, kita mengenalinya sebagai LSTM mundur, lalu mengelompokkan LSTM maju dan mundur bersama-sama. Ini adalah upaya masa depan. Saat ini, perintah ini akan membuat dua operasi UnidirectionalSequenceLSTM di model TensorFlow Lite.

Contoh konversi LSTM buatan pengguna

TensorFlow Lite juga menyediakan cara untuk mengonversi implementasi LSTM yang ditentukan pengguna. Di sini, kita menggunakan LSTM Lingvo sebagai contoh cara penerapannya. Untuk mengetahui detailnya, lihat antarmuka lingvo.LSTMCellSimple dan logika konversi di sini. Kami juga menyediakan contoh untuk definisi LSTM Lingvo lainnya di antarmuka lingvo.LayerNormalizedLSTMCellSimple dan logika konversinya di sini.

"Bawa TensorFlow RNN Anda sendiri" ke TensorFlow Lite

Jika antarmuka RNN pengguna berbeda dengan antarmuka standar yang didukung, ada beberapa opsi:

Opsi 1: Tulis kode adaptor di TensorFlow Python untuk menyesuaikan antarmuka RNN ke antarmuka Keras RNN. Hal ini berarti tf.function dengan anotasi tf_implementations pada fungsi antarmuka RNN yang dihasilkan yang identik dengan fungsi yang dihasilkan oleh lapisan Keras LSTM. Setelah ini, API konversi yang sama yang digunakan untuk Keras LSTM akan berfungsi.

Opsi 2: Jika hal di atas tidak memungkinkan (misalnya Keras LSTM tidak memiliki beberapa fungsi yang saat ini diekspos oleh op LSTM op seperti normalisasi lapisan TensorFlow Lite), maka perluas konverter TensorFlow Lite dengan menulis kode konversi kustom dan colokkan ke fungsi prepare-composite-MLIR-di sini. Antarmuka fungsi harus diperlakukan seperti kontrak API dan harus berisi argumen yang diperlukan untuk dikonversi ke operasi TensorFlow Lite LSTM gabungan - yaitu input, bias, bobot, proyeksi, normalisasi lapisan, dll. Sebaiknya tensor yang diteruskan sebagai argumen ke fungsi ini memiliki peringkat yang diketahui (yaitu RankedTensorType di MLIR). Hal ini mempermudah penulisan kode konversi yang dapat mengasumsikan tensor ini sebagai RankedTensorType dan membantu mengubahnya menjadi tensor peringkat yang sesuai dengan operand operator TensorFlow Lite gabungan.

Contoh lengkap dari alur konversi tersebut adalah konversi LSTMCellSimple ke TensorFlow Lite dari Lingvo.

LSTMCellSimple di Lingvo didefinisikan di sini. Model yang dilatih dengan sel LSTM ini dapat dikonversi ke TensorFlow Lite sebagai berikut:

  1. Gabungkan semua penggunaan LSTMCellSimple dalam tf.function dengan anotasi tf_implementations yang diberi label demikian (misalnya, lingvo.LSTMCellSimple akan menjadi nama anotasi yang baik di sini). Pastikan tf.function yang dihasilkan cocok dengan antarmuka fungsi yang diharapkan dalam kode konversi. Ini adalah kontrak antara penulis model yang menambahkan anotasi dan kode konversi.

  2. Perpanjang pass fungsi prepare-composite untuk menghubungkan operasi gabungan kustom ke TensorFlow Lite yang menggabungkan konversi operasi LSTM. Lihat kode konversi LSTMCellSimple.

    Kontrak konversi:

  3. Tensor Weight dan projection ditransposisikan.

  4. {input, recurrent} ke {cell, input gate, remember gate, output gate} diekstrak dengan mengiris tensor bobot yang ditransposisikan.

  5. {bias} ke {cell, input gate, remember gate, output gate} diekstrak dengan memotong tensor bias.

  6. Proyeksi diekstrak dengan memotong tensor proyeksi yang ditransposisikan.

  7. Konversi serupa juga ditulis untuk LayerNormalizedLSTMCellSimple.

  8. Infrastruktur konversi TensorFlow Lite lainnya, termasuk semua pass MLIR yang ditentukan serta ekspor akhir ke flatbuffer TensorFlow Lite dapat digunakan kembali.

Masalah umum/batasan

  1. Saat ini hanya ada dukungan untuk mengonversi Keras LSTM stateless (perilaku default di Keras). Konversi Stateful Keras LSTM adalah pekerjaan di masa mendatang.
  2. Anda masih dapat membuat model lapisan Keras LSTM stateful menggunakan lapisan Keras LSTM stateless yang mendasarinya dan mengelola status secara eksplisit dalam program pengguna. Program TensorFlow tersebut masih dapat dikonversi ke TensorFlow Lite menggunakan fitur yang dijelaskan di sini.
  3. LSTM dua arah saat ini dimodelkan sebagai dua operasi UnidirectionalSequenceLSTM di TensorFlow Lite. Ini akan diganti dengan operasi BidirectionalSequenceLSTM tunggal.