Konversi TensorFlow RNN ke TensorFlow Lite

Ringkasan

TensorFlow Lite mendukung konversi model TensorFlow RNN ke operasi LSTM gabungan TensorFlow Lite. Operasi gabungan ada untuk memaksimalkan kinerja 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 ada dua:

  1. Berikan dukungan asli untuk API RNN TensorFlow standar seperti Keras LSTM. Ini adalah opsi yang disarankan.
  2. Menyediakan antarmuka ke dalam infrastruktur konversi untuk implementasi RNN yang ditentukan pengguna agar dapat dipasang dan dikonversi ke TensorFlow Lite. Kami memberikan beberapa contoh luar biasa dari konversi tersebut menggunakan antarmuka LSTMCellSimple dan LayerNormalizedLSTMCellSimple RNN lingvo.

API Konverter

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

Fungsionalitas konversi ini tersedia saat mengonversi ke TensorFlow Lite melalui SavedModel atau dari model Keras secara langsung. Lihat contoh penggunaan.

Dari model yang disimpan

# 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 hingga TensorFlow Lite Colab mengilustrasikan penggunaan ujung ke ujung dengan interpreter TensorFlow Lite.

API TensorFlow RNN didukung

Kami mendukung konversi langsung Keras LSTM ke TensorFlow Lite. Untuk detail tentang cara kerjanya silakan merujuk ke antarmuka Keras LSTM dan logika konversi di sini .

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

  1. Dimensi 0 dari tensor masukan adalah ukuran batch.
  2. Dimensi 0 dari tensor recurrent_weight adalah jumlah keluaran.
  3. Tensor bobot dan recurrent_kernel dialihkan urutannya.
  4. Bobot yang ditransposisikan, kernel_rekuren yang ditransposisikan, dan tensor bias dibagi menjadi 4 tensor berukuran sama sepanjang dimensi 0. Ini sesuai dengan gerbang input, gerbang lupa, sel, dan gerbang keluaran .

Varian Keras LSTM

Jurusan waktu

Pengguna dapat memilih waktu utama atau tanpa waktu utama. Keras LSTM menambahkan atribut waktu-utama dalam atribut fungsi def. Untuk LSTM urutan searah, kita cukup 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 contoh di sini . Setelah kita melihat atribut go_backward, kita mengenalinya sebagai LSTM mundur, lalu kita mengelompokkan LSTM maju & mundur menjadi satu. Ini adalah pekerjaan masa depan. Saat ini, hal ini menciptakan dua operasi UnidirectSequenceLSTM dalam model TensorFlow Lite.

Contoh konversi LSTM yang ditentukan pengguna

TensorFlow Lite juga menyediakan cara untuk mengonversi implementasi LSTM yang ditentukan pengguna. Di sini kami menggunakan LSTM Lingvo sebagai contoh penerapannya. Untuk detailnya silakan merujuk ke antarmuka lingvo.LSTMCellSimple dan logika konversi di sini . Kami juga memberikan contoh definisi LSTM Lingvo lainnya di antarmuka lingvo.LayerNormalizedLSTMCellSimple dan logika konversinya di sini .

“Bawakan TensorFlow RNN Anda sendiri” ke TensorFlow Lite

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

Opsi 1: Tulis kode adaptor di TensorFlow python untuk mengadaptasi antarmuka RNN ke antarmuka Keras RNN. Ini berarti tf.function dengan anotasi tf_implements pada fungsi antarmuka RNN yang dihasilkan yang identik dengan 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 normalisasi lapisan seperti operasi LSTM TensorFlow Lite), perluas konverter TensorFlow Lite dengan menulis kode konversi khusus dan sambungkan ke persiapan -fungsi-komposit MLIR-lulus di sini . Antarmuka fungsi harus diperlakukan seperti kontrak API dan harus berisi argumen yang diperlukan untuk mengonversi operasi TensorFlow Lite LSTM yang digabungkan - yaitu input, bias, bobot, proyeksi, normalisasi lapisan, dll. Sebaiknya tensor diteruskan sebagai argumen untuk ini berfungsi untuk mengetahui peringkat (yaitu PeringkatTensorType di MLIR). Hal ini mempermudah penulisan kode konversi yang dapat menganggap tensor ini sebagai PeringkatTensorType dan membantu mengubahnya menjadi tensor peringkat yang sesuai dengan operan operator TensorFlow Lite yang digabungkan.

Contoh lengkap alur konversi tersebut adalah konversi LSTMCellSimple ke TensorFlow Lite 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_implements yang diberi label seperti itu (misalnya lingvo.LSTMCellSimple akan menjadi nama anotasi yang bagus di sini). Pastikan tf.function yang dihasilkan cocok dengan antarmuka fungsi yang diharapkan dalam kode konversi. Ini adalah kontrak antara pembuat model yang menambahkan anotasi dan kode konversi.
  2. Perluas pass fungsi persiapan-komposit untuk menyambungkan operasi komposit khusus ke konversi operasi LSTM gabungan TensorFlow Lite. Lihat kode konversi LSTMCellSimple .

    Kontrak konversi:

  3. Tensor berat dan proyeksi dialihkan posisinya.

  4. {input, recurrent} ke {cell, input gate, forget gate, output gate} diekstraksi dengan memotong tensor bobot yang dialihkan.

  5. {bias} ke {sel, gerbang masukan, gerbang lupa, gerbang keluaran} diekstraksi dengan memotong tensor bias.

  6. Proyeksi diekstraksi dengan memotong tensor proyeksi yang dialihkan.

  7. Konversi serupa 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/keterbatasan yang diketahui

  1. Saat ini hanya ada dukungan untuk mengonversi Keras LSTM tanpa kewarganegaraan (perilaku default di Keras). Konversi Stateful Keras LSTM adalah pekerjaan masa depan.
  2. Lapisan Keras LSTM stateful masih dapat dimodelkan 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 UnidirectSequenceLSTM di TensorFlow Lite. Ini akan diganti dengan satu operasi BiDirectionalSequenceLSTM.