Konversi RNN TensorFlow ke LiteRT

Ringkasan

LiteRT mendukung konversi model RNN TensorFlow ke operasi LSTM gabungan LiteRT. Operasi gabungan ada untuk memaksimalkan performa implementasi kernel yang mendasarinya, serta menyediakan antarmuka tingkat yang lebih tinggi untuk menentukan transformasi kompleks seperti kuantisasi.

Karena ada banyak varian RNN API di TensorFlow, pendekatan kami adalah dua arah:

  1. Menyediakan dukungan native untuk API RNN TensorFlow standar seperti Keras LSTM. Ini adalah opsi yang direkomendasikan.
  2. Menyediakan antarmuka ke infrastruktur konversi untuk penerapan RNN yang ditentukan pengguna agar dapat dihubungkan dan dikonversi ke LiteRT. Kami menyediakan beberapa contoh konversi langsung seperti itu menggunakan LSTMCellSimple dan LayerNormalizedLSTMCellSimple antarmuka RNN lingvo.

Converter API

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

Fungsi konversi ini tersedia saat mengonversi ke LiteRT melalui SavedModel 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 LiteRT Colab mengilustrasikan penggunaan end-to-end dengan penafsir LiteRT.

API RNN TensorFlow yang didukung

Kami mendukung konversi Keras LSTM ke LiteRT secara langsung. Untuk mengetahui detail cara kerjanya, lihat antarmuka Keras LSTM dan logika konversi di sini.

Yang juga penting adalah menyoroti kontrak LSTM LiteRT sehubungan dengan definisi operasi Keras:

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

Varian LSTM Keras

Waktu utama

Pengguna dapat memilih time-major atau tidak ada time-major. LSTM Keras menambahkan atribut time-major dalam atribut def fungsi. Untuk LSTM urutan satu arah, kita dapat memetakan ke atribut time major unidirecional_sequence_lstm.

LSTM Dua Arah (BiDirectional LSTM)

LSTM dua arah dapat diimplementasikan dengan dua lapisan LSTM Keras, satu untuk maju dan satu untuk mundur, lihat contoh di sini. Setelah melihat atribut go_backward, kami akan mengenalinya sebagai LSTM mundur, lalu mengelompokkan LSTM maju & mundur bersama-sama. Ini adalah pekerjaan mendatang. Saat ini, hal ini akan membuat dua operasi UnidirectionalSequenceLSTM dalam model LiteRT.

Contoh konversi LSTM yang ditentukan pengguna

LiteRT 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 memberikan contoh untuk definisi LSTM Lingvo lainnya di antarmuka lingvo.LayerNormalizedLSTMCellSimple dan logika konversinya di sini.

“Bring your own TensorFlow RNN” ke LiteRT

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

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

Opsi 2: Jika hal di atas tidak memungkinkan (misalnya, LSTM Keras tidak memiliki beberapa fungsi yang saat ini diekspos oleh operasi LSTM gabungan LiteRT seperti normalisasi lapisan), perluas konverter LiteRT dengan menulis kode konversi kustom dan hubungkan ke MLIR-pass prepare-composite-functions di sini. Antarmuka fungsi harus diperlakukan seperti kontrak API dan harus berisi argumen yang diperlukan untuk mengonversi ke operasi LSTM LiteRT 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 akan mempermudah penulisan kode konversi yang dapat mengasumsikan tensor ini sebagai RankedTensorType dan membantu mentransformasinya menjadi tensor berperingkat yang sesuai dengan operan operator LiteRT gabungan.

Contoh lengkap alur konversi tersebut adalah konversi LSTMCellSimple Lingvo ke LiteRT.

LSTMCellSimple di Lingvo ditentukan di sini. Model yang dilatih dengan sel LSTM ini dapat dikonversi ke LiteRT sebagai berikut:

  1. Bungkus semua penggunaan LSTMCellSimple dalam tf.function dengan anotasi tf_implements yang diberi label seperti itu (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. Perluas pass prepare-composite-functions untuk memasukkan operasi komposit kustom ke konversi operasi LSTM gabungan LiteRT. Lihat kode konversi LSTMCellSimple.

    Kontrak konversi:

  3. Tensor bobot dan proyeksi ditransposisikan.

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

  5. {bias} ke {cell, input gate, forget gate, output gate} diekstrak dengan mengiris tensor bias.

  6. Proyeksi diekstrak dengan mengiris tensor proyeksi yang ditransposisikan.

  7. Konversi serupa ditulis untuk LayerNormalizedLSTMCellSimple.

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

Masalah/batasan umum

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