TensorFlow RNN 轉換至 TensorFlow Lite

總覽

TensorFlow Lite 支援將 TensorFlow RNN 模型轉換成 TensorFlow Lite 的整合式 LSTM 運算。融合作業是為了最大化基礎核心實作的效能，並提供更高層級的介面來定義量化等複雜轉換作業。

由於 TensorFlow 中有許多 RNN API 變體，我們採取的方法分為兩種：

  1. 提供標準 TensorFlow RNN API 的原生支援，例如 Keras LSTM。這是我們建議的選項。
  2. 將「介面」介面提供給轉換基礎架構，以便 使用者定義的 RNN 實作項目進行插入，並轉換為 TensorFlow Lite。我們使用 lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 介面提供幾個立即可用的轉換範例。

Converter API

這項功能是 TensorFlow 2.3 版本的一部分。您也可以透過 tf-nightly 管道或從頭存取。

透過 MRAID 或直接從 Keras 模型轉換為 TensorFlow Lite 時，可以使用這項轉換功能。查看使用範例。

來自已儲存的模型

# 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 轉換為 TensorFlow Lite。如要進一步瞭解相關運作方式，請參閱 Keras LSTM 介面 和這裡的轉換邏輯。

此外，還要特別強調 TensorFlow Lite 與 Keras 作業定義的 TensorFlow Lite 合約：

  1. 輸入張量的維度 0 為批量。
  2. recurrent_weight 張量的維度 0 是輸出數量。
  3. 系統會轉移 weight 和 recurrent_kernel 張量。
  4. 轉置的權重、轉置的 recurrent_kernel 和 bias 張量在維度 0 上分為 4 個同等大小的張量。這些控制項會對應至輸入閘、清除閘門、儲存格和輸出閘。

Keras LSTM 子類

時間主要

使用者可以選擇「主要」或「不使用時間」。Keras LSTM 會在函式 def 屬性中新增時間主要屬性。針對單向序列 LSTM，我們只需對應到 unidirecional_sequence_lstm 的時間主要屬性。

BiDirectional LSTM

Bidirectional LSTM 可透過兩個 Keras LSTM 層實作，一個用於正向和反向，請參閱這裡的範例。當我們看到 go_backward 屬性時，我們會將其視為回溯 LSTM，然後再區分前和向後 LSTM。這是未來的措施。目前，這會在 TensorFlow Lite 模型中建立兩個 UnidirectionalSequenceLSTM 作業。

使用者定義的 LSTM 轉換範例

TensorFlow Lite 也提供轉換使用者定義的 LSTM 實作項目。我們在此以 Lingvo 的 LSTM 做為實作範例。詳情請參閱 lingvo.LSTMCellSimple 介面和這裡的轉換邏輯。我們也提供有關 Lingvo 的其他 LSTM 定義範例，請參閱 lingvo.LayerNormalizedLSTMCellSimple 介面，以及此處的轉換邏輯。

TensorFlow Lite 的「自帶 TensorFlow RNN」

如果使用者的 RNN 介面與標準支援的介面不同，可以採取以下幾種做法：

選項 1：在 TensorFlow Python 中編寫轉接器程式碼，將 RNN 介面調整為 Keras RNN 介面。這表示一個在產生的 RNN 介面函式中使用 tf_implementations 註解的 tf. 函式，與 Keras LSTM 層產生的註解相同。完成後，用於 Keras LSTM 的轉換 API 就能正常運作。

選項 2：如果無法執行上述操作 (例如 Keras LSTM 缺少部分由 TensorFlow Lite 融合的 LSTM 運算 (例如層正規化) 公開的功能)，請編寫自訂轉換程式碼並將其插入這裡的 prepare 複合函式 MLIR 傳遞功能來擴充 TensorFlow Lite 轉換工具。函式的介面應視為 API 合約，且應包含轉換時融合的 TensorFlow Lite LSTM 作業所需的引數，例如輸入、偏誤、權重、投影、層正規化等。建議以引數形式傳遞的張量要達到已知排名 (即 MLIR 中的排名)。如此一來，您就能更輕鬆地編寫轉換程式碼，將這些張量假設為 RankedTensorType，並協助將這類張量轉換為與整合式 TensorFlow Lite 運算元對應的排序張量。

從 Lingvo 的 LSTMCellSimple 轉換至 TensorFlow Lite。

Lingvo 中的 LSTMCellSimple 請參閱此處的定義。使用這個 LSTM 儲存格訓練的模型可以轉換為 TensorFlow Lite，如下所示：

  1. 將所有 LSTMCellSimple 所有相關用途納入 tf.function，並加上 tf_implementations 註解，例如類似這樣的標籤 (例如 lingvo.LSTMCellSimple)，即是很好的註解名稱。確認產生的 tf.function 與轉換程式碼中預期的函式介面相符。這是模型作者在新增註解和轉換程式碼之間的合約。

  2. 將 prepare-複合函式傳遞的範圍延伸至 TensorFlow Lite 融合的 LSTM 運算轉換。請參閱 LSTMCellSimple 轉換程式碼。

    轉換合約：

  3. 重量和投影張量會轉置。

  4. 系統會分割轉置的權重張量，藉此擷取 {input, recurrent} 至 {cell, input gate, forgot gate, outputgate}。

  5. 系統會複製自訂調整張量，藉此擷取 {bias} 至 {cell, input gate, forgot gate, output gate}。

  6. 對轉置的投影張量分割投影。

  7. 類似轉換是針對 LayerNormalizedLSTMCellSimple。

  8. TensorFlow Lite 轉換基礎架構的其餘部分包括定義的所有 MLIR 傳遞，以及匯出至 TensorFlow Lite 扁平緩衝區的最終結果。

已知問題/限制

  1. 目前只支援轉換無狀態 Keras LSTM (Keras 中的預設行為)。有狀態的 Keras LSTM 轉換是未來的工作。
  2. 不過，您還是可以使用基礎的無狀態 Keras LSTM 層，建立有狀態的 Keras LSTM 層模型，並在使用者程式中明確管理狀態。您仍然可以使用本文說明的功能，將這類 TensorFlow 程式轉換為 TensorFlow Lite。
  3. 目前在 TensorFlow Lite 中，雙向 LSTM 模擬為兩項 UnidirectionalSequenceLSTM 作業。這會替換為單一 BidirectionalSequenceLSTM 操作。