TensorFlow RNN を TensorFlow Lite に変換する

概要

TensorFlow Lite は、TensorFlow RNN モデルを TensorFlow Lite の融合 LSTM 演算に変換することをサポートしています。融合演算は、基盤となるカーネル実装のパフォーマンスを最大化するとともに、量子化などの複雑な変換を定義するための上位レベルのインターフェースを提供するものです。

TensorFlow には RNN API の多くのバリエーションがあるため、Google のアプローチは次の 2 つに分かれています。

  1. Keras LSTM などの標準の TensorFlow RNN API のネイティブ サポートを提供します。これが推奨のオプションです。
  2. ユーザー定義の RNN 実装をプラグインして TensorFlow Lite に変換する変換インフラストラクチャへのインターフェースを提供します。Google では、lingvo の LSTMCellSimple および LayerNormalizedLSTMCellSimple RNN インターフェースを使用した、すぐに使える変換例をいくつか紹介します。

Converter API

この機能は TensorFlow 2.3 リリースの一部です。tf-nightly の pip または head からも確認できます。

この変換機能は、SavedModel 経由、または 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 RNNs API

Google は、Keras LSTM から TensorFlow Lite への変換をすぐにサポートしています。この仕組みの詳細については、Keras LSTM インターフェース と、こちらの変換ロジックを参照してください。

また、Keras オペレーションの定義に関する TensorFlow Lite の LSTM 契約を強調することも重要です。

  1. input テンソルのディメンション 0 はバッチサイズです。
  2. recurrent_weight テンソルの次元 0 は出力の数です。
  3. weight テンソルと recurrent_kernel テンソルが転置されます。
  4. 転置した重み、転置した recurrent_kernel、bias のテンソルは、次元 0 に沿って 4 つの同じサイズのテンソルに分割されます。これらは、入力ゲート、フォーゲット ゲート、セル、出力ゲートに対応します。

Keras LSTM バリアント

メジャー

ユーザーは「Time-major」または「Time-major」を選択できます。Keras LSTM では、関数 def 属性に時間優先属性が追加されます。単方向シーケンス LSTM の場合は、unidirecional_sequence_lstm の時刻メジャー属性に単純にマッピングできます。

双方向 LSTM

双方向 LSTM は、2 つの Keras LSTM レイヤ（1 つはフォワード用、もう 1 つはバック用）を使用して実装できます。例については、こちらをご覧ください。go_backward 属性では、後方 LSTM として認識し、前方 LSTM と後方 LSTM をグループ化しています。これは将来の取り組みです。現在、これにより TensorFlow Lite モデルに 2 つの UnidirectionalSequenceLSTM オペレーションが作成されます。

ユーザー定義の LSTM 変換の例

TensorFlow Lite では、ユーザー定義の LSTM 実装を変換することもできます。ここでは、その実装方法の例として Lingvo の LSTM を使用します。詳しくは、lingvo.LSTMCellSimple インターフェースとこちらの変換ロジックをご覧ください。また、lingvo.LayerNormalizedLSTMCellSimple インターフェースにある別の Lingvo の LSTM 定義の例と、こちらの変換ロジックの例も紹介します。

「独自の TensorFlow RNN を TensorFlow Lite で使用する」

ユーザーの RNN インターフェースが標準でサポートされているものと異なる場合は、次の 2 つの方法があります。

オプション 1: TensorFlow Python でアダプタコードを記述して、RNN インターフェースを Keras RNN インターフェースに適応させます。つまり、生成された RNN インターフェースの関数に、Keras LSTM レイヤが生成したものと同一の tf_Implements アノテーションを持つ tf.function が存在します。その後は、Keras LSTM に使用するのと同じコンバージョン API が機能します。

オプション 2: 上記が不可能な場合（例: Keras LSTM に、レイヤ正規化など TensorFlow Lite の融合 LSTM 演算によって現在公開されている機能が不足している場合など）、カスタム変換コードを記述して TensorFlow Lite コンバータを拡張し、こちらの prepare-composite-functions MLIR-pass に埋め込みます。関数のインターフェースは API コントラクトのように扱う必要があり、融合 TensorFlow Lite LSTM 演算に変換するために必要な引数（入力、バイアス、重み、射影、レイヤ正規化など）が含まれている必要があります。この関数に引数として渡されるテンソルには、既知の階数（MLIR の RankedTensorType）があることが推奨されます。これにより、これらのテンソルを RankedTensorType として想定できる変換コードを記述し、融合 TensorFlow Lite 演算子のオペランドに対応するランク付けされたテンソルに変換できます。

このような変換フローの完全な例は、Lingvo の LSTMCellSimple から TensorFlow Lite への変換です。

Lingvo の LSTMCellSimple は、こちらで定義されています。この LSTM セルでトレーニングされたモデルは、次のように TensorFlow Lite に変換できます。

  1. tf.function 内で LSTMCellSimple が使用されている場合は、そのようにラベル付けされた tf_Implements アノテーションを使用してラップします（たとえば、ここでは lingvo.LSTMCellSimple が適切なアノテーション名になります）。生成される tf.function が、変換コードで想定される関数のインターフェースと一致することを確認します。これは、モデル作成者がアノテーションを追加することと、変換コードを追加することの間の契約です。

  2. prepare-composite-functions パスを拡張して、カスタム複合演算を TensorFlow Lite 融合 LSTM 演算変換に接続します。LSTMCellSimple 変換コードをご覧ください。

    変換契約:

  3. 重みテンソルと射影テンソルが転置されます。

  4. {input, recurrent} から {cell, input gate, never gate, output gate} は、転置された重みテンソルをスライスすることで抽出されます。

  5. {bias} から {cell, input gate, へ移動 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 が 2 つの UnidirectionalSequenceLSTM 演算としてモデル化されています。これは、単一の BidirectionalSequenceLSTM 演算に置き換えられます。