總覽
TensorFlow Lite 支援將 TensorFlow RNN 模型轉換為 TensorFlow Lite 的融合 LSTM 運算。融合運算旨在盡可能提升底層核心實作的效能，並提供更高等級的介面，以定義量化等複雜轉換。
由於 TensorFlow 中有許多 RNN API 變體，因此我們採取了雙管齊下的方法：
- 為標準 TensorFlow RNN API (例如 Keras LSTM) 提供原生支援。這是建議選項。
- 為使用者定義的 RNN 實作提供介面連入轉換基礎架構，以便外掛並轉換為 TensorFlow Lite。我們提供幾個現成的範例，說明如何使用 Lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 介面進行此類轉換。
轉換器 API
此功能是 TensorFlow 2.3 版本的一部分。也可透過 tf-nightly pip 或從主分支取得。
透過 SavedModel 或直接從 Keras 模型轉換為 TensorFlow Lite 時，此轉換功能即可使用。請參閱範例用法。
從 SavedModel
# 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()
從 Keras 模型
# build a Keras model
keras_model = build_keras_lstm(...)
# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()
範例
Keras LSTM 轉 TensorFlow Lite Colab 說明搭配 TensorFlow Lite 解譯器的端對端用法。
支援的 TensorFlow RNN API
Keras LSTM 轉換 (建議)
我們支援將 Keras LSTM 開箱即用轉換為 TensorFlow Lite。如需瞭解其運作方式詳情，請參閱 Keras LSTM 介面 以及此處的轉換邏輯。
同樣重要的是，要強調 TensorFlow Lite 的 LSTM 合約與 Keras 運算定義的關係
- 輸入張量的維度 0 是批次大小。
- recurrent_weight 張量的維度 0 是輸出數量。
- weight 和 recurrent_kernel 張量已轉置。
- 轉置的 weight、轉置的 recurrent_kernel 和 bias 張量沿著維度 0 分割成 4 個大小相等的張量。這些張量對應於輸入閘、遺忘閘、儲存格和輸出閘。
Keras LSTM 變體
時間主軸
使用者可以選擇時間主軸或非時間主軸。Keras LSTM 在函式定義屬性中新增了時間主軸屬性。對於單向序列 LSTM，我們可以簡單地對應至 unidirecional_sequence_lstm 的 時間主軸屬性。
雙向 LSTM
雙向 LSTM 可以使用兩個 Keras LSTM 層實作，一個用於正向，另一個用於反向，請參閱此處的範例。一旦我們看到 go_backward 屬性，就會將其識別為反向 LSTM，然後將正向和反向 LSTM 分組在一起。這是未來工作。目前，這會在 TensorFlow Lite 模型中建立兩個 UnidirectionalSequenceLSTM 運算。
使用者定義的 LSTM 轉換範例
TensorFlow Lite 也提供一種轉換使用者定義的 LSTM 實作的方法。在這裡，我們以 Lingvo 的 LSTM 為例，說明如何實作。如需瞭解詳情，請參閱 lingvo.LSTMCellSimple 介面和此處的轉換邏輯。我們也在 lingvo.LayerNormalizedLSTMCellSimple 介面及其此處的轉換邏輯中，提供另一個 Lingvo LSTM 定義的範例。
「將您自己的 TensorFlow RNN 帶入」TensorFlow Lite
如果使用者的 RNN 介面與標準支援的介面不同，則有幾個選項：
選項 1：在 TensorFlow Python 中編寫轉接器程式碼，以將 RNN 介面調整為 Keras RNN 介面。這表示在產生的 RNN 介面函式上使用具有 tf_implements 註解的 tf.function，該函式與 Keras LSTM 層產生的函式相同。完成此操作後，用於 Keras LSTM 的相同轉換 API 即可運作。
選項 2：如果上述方法不可行 (例如，Keras LSTM 缺少 TensorFlow Lite 的融合 LSTM 運算 (例如層正規化) 目前公開的某些功能)，則擴充 TensorFlow Lite 轉換器，方法是編寫自訂轉換程式碼並將其外掛到 prepare-composite-functions MLIR 傳遞 此處。函式的介面應視為 API 合約，且應包含轉換為融合 TensorFlow Lite LSTM 運算 (即輸入、偏差、權重、投影、層正規化等) 所需的引數。最好將作為此函式引數傳遞的張量設為已知等級 (即 MLIR 中的 RankedTensorType)。這讓編寫轉換程式碼變得更容易，因為程式碼可以將這些張量視為 RankedTensorType，並協助將其轉換為對應於融合 TensorFlow Lite 運算子運算元的已排名張量。
此類轉換流程的完整範例是 Lingvo 的 LSTMCellSimple 轉 TensorFlow Lite 轉換。
Lingvo 中的 LSTMCellSimple 定義於此處。使用此 LSTM 儲存格訓練的模型可以轉換為 TensorFlow Lite，方法如下：
- 將 LSTMCellSimple 的所有用法包裝在 tf.function 中，並加上標示為此類型的 tf_implements 註解 (例如，lingvo.LSTMCellSimple 在此處會是良好的註解名稱)。確保產生的 tf.function 符合轉換程式碼中預期的函式介面。這是新增註解的模型作者與轉換程式碼之間的合約。
擴充 prepare-composite-functions 傳遞，以將自訂複合運算子外掛到 TensorFlow Lite 融合 LSTM 運算子轉換。請參閱 LSTMCellSimple 轉換程式碼。
轉換合約
Weight 和 projection 張量已轉置。
{輸入、循環} 至 {儲存格、輸入閘、遺忘閘、輸出閘} 是透過分割轉置的 weight 張量來擷取。
{偏差} 至 {儲存格、輸入閘、遺忘閘、輸出閘} 是透過分割偏差張量來擷取。
Projection 是透過分割轉置的 projection 張量來擷取。
LayerNormalizedLSTMCellSimple 也編寫了類似的轉換，位置在此處。
TensorFlow Lite 轉換基礎架構的其餘部分，包括所有定義的 MLIR 傳遞以及最終匯出至 TensorFlow Lite FlatBuffer，都可以重複使用。
已知問題/限制
- 目前僅支援轉換無狀態 Keras LSTM (Keras 中的預設行為)。有狀態 Keras LSTM 轉換是未來工作。
- 仍然可以使用底層無狀態 Keras LSTM 層對有狀態 Keras LSTM 層進行建模，並在使用者程式中明確管理狀態。此處所述的功能仍然可用於將此類 TensorFlow 程式轉換為 TensorFlow Lite。
- 雙向 LSTM 目前在 TensorFlow Lite 中建模為兩個 UnidirectionalSequenceLSTM 運算。這將替換為單一 BidirectionalSequenceLSTM 運算。