TensorFlow RNN 轉換為 LiteRT

總覽

LiteRT 支援將 TensorFlow RNN 模型轉換為 LiteRT 的融合 LSTM 運算。融合作業可盡量提升基礎核心實作的效能，並提供較高層級的介面，定義量化等複雜轉換。

由於 TensorFlow 中有許多 RNN API 變體，我們的做法分為兩部分：

  1. 原生支援標準 TensorFlow RNN API，例如 Keras LSTM。 這是我們建議的選項。
  2. 提供介面 ，供使用者定義的 RNN 實作項目插入轉換基礎架構，並轉換為 LiteRT。我們提供幾個現成的範例，說明如何使用 lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 介面進行這類轉換。

Converter API

這項功能是 TensorFlow 2.3 版本的一部分。您也可以透過 tf-nightly pip 或從主幹取得。

透過 SavedModel 或直接從 Keras 模型轉換為 LiteRT 時，即可使用這項轉換功能。請參閱使用範例。

從已儲存的模特兒

# 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 至 LiteRT Colab 說明如何搭配 LiteRT 解譯器使用端對端。

支援的 TensorFlow RNN API

我們支援將 Keras LSTM 轉換為 LiteRT，如要瞭解相關運作方式，請參閱 Keras LSTM 介面 和這裡的轉換邏輯。

此外，請務必根據 Keras 作業定義，強調 LiteRT 的 LSTM 合約：

  1. input 張量的維度 0 是批次大小。
  2. recurrent_weight 張量的維度 0 是輸出數量。
  3. weight 和 recurrent_kernel 張量會轉置。
  4. 轉置權重、轉置 recurrent_kernel 和 bias 張量會沿著維度 0 分割成 4 個大小相等的張量。這些分別對應至輸入閘、遺忘閘、儲存格和輸出閘。

Keras LSTM 變體

時間主軸

使用者可以選擇主修時間相關科目，也可以不選擇。Keras LSTM 會在函式定義屬性中新增時間主要屬性。如果是單向序列 LSTM，我們可以簡單地對應至 unidirecional_sequence_lstm 的時間主要屬性。

雙向 LSTM

雙向 LSTM 可透過兩個 Keras LSTM 層實作，一個用於正向，另一個用於反向，請參閱這裡的範例。一旦看到 go_backward 屬性，系統就會將其視為反向 LSTM，然後將正向和反向 LSTM 分組。這是未來的工作。目前，這會在 LiteRT 模型中建立兩個 UnidirectionalSequenceLSTM 作業。

使用者定義的 LSTM 轉換範例

LiteRT 也提供轉換使用者定義 LSTM 實作的方法。我們以 Lingvo 的 LSTM 為例，說明如何實作這項功能。詳情請參閱 lingvo.LSTMCellSimple 介面，以及這裡的轉換邏輯。我們也提供 Lingvo 的另一個 LSTM 定義範例，請參閱 lingvo.LayerNormalizedLSTMCellSimple 介面和這裡的轉換邏輯。

「將您自己的 TensorFlow RNN」帶入 LiteRT

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

選項 1：在 TensorFlow Python 中編寫轉接器程式碼，將 RNN 介面調整為 Keras RNN 介面。也就是說，tf.function 具有 tf_implements 註解，位於產生的 RNN 介面函式上，與 Keras LSTM 層產生的函式相同。完成後，您就能使用 Keras LSTM 適用的轉換 API。

選項 2：如果無法使用上述方法 (例如 Keras LSTM 缺少 LiteRT 融合 LSTM 作業目前公開的某些功能，如層級正規化)，請編寫自訂轉換程式碼，並將其插入 prepare-composite-functions MLIR 傳遞這裡，藉此擴充 LiteRT 轉換器。 函式的介面應視為 API 合約，且應包含轉換為融合 LiteRT LSTM 作業所需的引數，也就是輸入、偏差、權重、投影、層級正規化等。傳遞為這個函式引數的張量最好具有已知等級 (即 MLIR 中的 RankedTensorType)。這樣一來，撰寫轉換程式碼時，就能輕鬆將這些張量視為 RankedTensorType，並轉換為對應於融合 LiteRT 運算元運算元的排序張量。

這類轉換流程的完整範例是 Lingvo 的 LSTMCellSimple 到 LiteRT 轉換。

Lingvo 中的 LSTMCellSimple 定義位於這裡。使用這個 LSTM 儲存格訓練的模型可轉換為 LiteRT，方法如下：

  1. 將 LSTMCellSimple 的所有用法包裝在 tf.function 中，並使用標示為這類函式的 tf_implements 註解 (例如 lingvo.LSTMCellSimple 在這裡會是合適的註解名稱)。請確認產生的 tf.function 符合轉換程式碼中預期的函式介面。這是模型作者 (新增註解) 與轉換程式碼之間的合約。

  2. 擴充 prepare-composite-functions 傳遞，插入自訂複合作業，以轉換為 LiteRT 融合 LSTM 作業。請參閱 LSTMCellSimple 轉換程式碼。

    轉換合約：

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

  4. 轉置權重張量會經過切片處理，擷取 {input, recurrent} 至 {cell, input gate, forget gate, output gate}。

  5. 偏誤張量會經過切片處理，擷取 {bias} 至 {cell、input gate、forget gate、output gate}。

  6. 轉置投影張量經過切片後，即可擷取投影。

  7. 類似的轉換會寫入 LayerNormalizedLSTMCellSimple。

  8. LiteRT 轉換基礎架構的其餘部分 (包括所有定義的 MLIR 傳遞，以及最終匯出至 LiteRT flatbuffer) 都可以重複使用。

已知問題/限制

  1. 目前僅支援轉換無狀態的 Keras LSTM (Keras 中的預設行為)。有狀態的 Keras LSTM 轉換是未來的工作。
  2. 您仍可使用基礎無狀態 Keras LSTM 層，並在使用者程式中明確管理狀態，藉此模擬有狀態的 Keras LSTM 層。這類 TensorFlow 程式仍可使用本文所述功能轉換為 LiteRT。
  3. 在 LiteRT 中，雙向 LSTM 目前會模擬為兩個 UnidirectionalSequenceLSTM 作業。這會由單一 BidirectionalSequenceLSTM 作業取代。