Konversi RNN TensorFlow ke LiteRT

Ringkasan

LiteRT mendukung konversi model RNN TensorFlow ke LiteRT gabungan operasi LSTM. Operasi gabungan tersedia untuk memaksimalkan performa implementasi {i>kernel<i} yang mendasarinya, serta menyediakan tingkat untuk mendefinisikan transformasi kompleks seperti kuantisasi.

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

  1. Memberikan dukungan native untuk TensorFlow RNN API standar seperti Keras LSTM. Opsi ini direkomendasikan.
  2. Berikan antarmuka ke dalam infrastruktur konversi untuk Penerapan RNN yang ditentukan pengguna untuk dicolokkan dan dikonversi menjadi LiteRT. Kami menyediakan beberapa contoh langsung konversi menggunakan konversi lingvo LSTMCellSimple dan LayerNormalizedLSTMCellSimple Antarmuka RNN.

API Konverter

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

Fungsi konversi ini tersedia saat melakukan konversi ke LiteRT 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 LiteRT Colab mengilustrasikan penggunaan menyeluruh dengan penafsir LiteRT.

TensorFlow RNN API yang didukung

Kami mendukung konversi siap pakai dari Keras LSTM ke LiteRT. Sebagai selengkapnya tentang cara kerjanya, lihat Antarmuka Keras LSTM dan ke logika konversi di sini.

Hal penting lainnya adalah menyoroti kontrak LSTM LiteRT sehubungan dengan ke 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 ditransposisi, recurrent_kernel, dan tensor bias yang ditransposisikan dibagi menjadi 4 tensor berukuran sama di sepanjang dimensi 0. Hal ini berkaitan dengan input gate, remember gate, sel, dan output gate.

Keras LSTM Variants

Waktu utama

Pengguna dapat memilih waktu sementara atau tanpa waktu sementara. Keras LSTM menambahkan di atribut def fungsi. Untuk LSTM urutan searah, kita cukup dipetakan ke unidirecional_sequence_lstm's atribut waktu utama.

LSTM BiDirectional

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

Contoh konversi LSTM yang ditentukan pengguna

LiteRT juga menyediakan cara untuk mengonversi LSTM yang ditentukan pengguna implementasi yang tepat. Di sini kita menggunakan LSTM Lingvo sebagai contoh bagaimana hal itu dapat diimplementasikan. Untuk mengetahui detailnya, lihat Antarmuka lingvo.LSTMCellSimple dan logika konversi di sini. Kami juga memberikan contoh untuk definisi LSTM Lingvo lainnya dalam Antarmuka lingvo.LayerNormalizedLSTMCellSimple dan logika konversinya di sini.

“Bawa TensorFlow RNN Anda sendiri” ke LiteRT

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

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

Opsi 2: Jika hal di atas tidak memungkinkan (misalnya Keras LSTM tidak menyertakan beberapa fungsionalitas yang saat ini diekspos oleh operasi LSTM gabungan dari LiteRT, lapisan normal), lalu memperluas konverter LiteRT dengan menulis kode konversi kustom dan memasukkannya ke dalam fungsi prepare-composite, MLIR-pass di sini. Antarmuka fungsi harus diperlakukan seperti kontrak API dan harus berisi argumen yang diperlukan untuk mengonversi ke LiteRT LSTM gabungan operasi - yaitu input, bias, bobot, proyeksi, normalisasi lapisan, dll. Ini lebih disarankan untuk tensor yang diteruskan sebagai argumen ke fungsi ini agar (yaitu RankedTensorType di MLIR). Hal ini membuatnya lebih mudah untuk menulis kode konversi yang dapat mengasumsikan tensor ini sebagai RankedTensorType dan membantu mengubahnya ke tensor peringkat yang sesuai dengan LiteRT gabungan operand operator.

Contoh lengkap dari alur konversi tersebut adalah LSTMCellSimple dari Lingvo untuk Konversi LiteRT.

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

  1. Menggabungkan semua penggunaan LSTMCellSimple dalam tf.function dengan tf_implementations yang diberi label seperti itu (mis. lingvo.LSTMCellSimple akan menjadi nama anotasi yang baik di sini). Memastikan tf.function yang dibuat cocok dengan antarmuka fungsi yang diharapkan dalam kode konversi. Ini adalah kontrak antara penulis model yang menambahkan anotasi dan kode konversi.
  2. Memperluas penerusan fungsi persiapan-komposit untuk menyambungkan operasi komposit kustom ke konversi operasi LSTM yang disatukan LiteRT. Lihat LSTMCellSimple kode konversi.

    Kontrak konversi:

  3. Tensor Bobot dan proyeksi ditransposisikan.

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

  5. {bias} ke {bias} adalah diekstrak dengan mengiris tensor bias.

  6. Proyeksi diekstrak dengan memotong tensor proyeksi yang ditransposisikan.

  7. Konversi serupa ditulis untuk LayerNormalizedLSTMCellSimple.

  8. Infrastruktur konversi LiteRT lainnya, termasuk semua Kartu MLIR ditentukan serta ekspor akhir ke {i>flatbuffer<i} LiteRT dapat digunakan kembali.

Masalah/batasan umum

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