TensorFlow RNN から TensorFlow Lite への変換

概要

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

TensorFlow には RNN API のバリエーションが多数あるため、私たちのアプローチは 2 つあります。

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

コンバーター 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 RNN API のサポート

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

また、Keras 操作定義に関する TensorFlow Lite の LSTM コントラクトを強調することも重要です。

  1. 入力テンソルの次元 0 はバッチ サイズです。
  2. recurrent_weightテンソルの次元 0 は出力の数です。
  3. Weight テンソルとrecurrent_kernelテンソルは転置されます。
  4. 転置された重み、転置された recurrent_kernel およびバイアステンソルは、次元 0 に沿って 4 つの等しいサイズのテンソルに分割されます。これらは、入力ゲート、忘却ゲート、セル、および出力ゲートに対応します。

Keras LSTM バリアント

時間専攻

ユーザーはタイム メジャーを選択するか、タイム メジャーを選択しないことを選択できます。 Keras LSTM は、関数 def 属性に時間優先属性を追加します。単方向シーケンス LSTM の場合、単純に unidiresional_sequence_lstm の時間メジャー属性にマッピングできます。

双方向LSTM

双方向 LSTM は、順方向と逆方向の 2 つの Keras LSTM レイヤーで実装できます。ここで例を参照してください。 go_backward 属性を確認すると、それが逆方向 LSTM であると認識し、順方向と逆方向 LSTM をグループ化します。これは今後の課題です。現在、これにより TensorFlow Lite モデルに 2 つの UnidirectionSequenceLSTM オペレーションが作成されます。

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

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

TensorFlow Lite への「独自の TensorFlow RNN の持ち込み」

ユーザーの 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 コンバータを拡張し、それを準備にプラグインします。 -composite-functions MLIR-passここに。関数のインターフェイスは API コントラクトのように扱われ、融合された TensorFlow Lite LSTM 操作に変換するために必要な引数 (入力、バイアス、重み、射影、レイヤー正規化など) を含む必要があります。これには、引数として渡されるテンソルが望ましいです。既知のランクを持つ関数 (つまり、MLIR の RankedTensorType)。これにより、これらのテンソルを RankedTensorType として想定できる変換コードの作成がはるかに簡単になり、融合された TensorFlow Lite オペレーターのオペランドに対応するランク付けされたテンソルに変換するのに役立ちます。

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

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

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

  2. prepare-composite-functions パスを拡張して、カスタム複合演算を TensorFlow Lite 融合 LSTM 演算変換にプラグインします。 LSTMCellSimple変換コードを参照してください。

    変換契約:

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

  4. {input、recurrent}から{cell、inputgate、forgetgate、outputgate}は、転置された重みテンソルをスライスすることによって抽出されます。

  5. {バイアス}から{セル、入力ゲート、忘却ゲート、出力ゲート}までは、バイアス テンソルをスライスすることによって抽出されます。

  6. 投影は、転置された投影テンソルをスライスすることによって抽出されます。

  7. 同様の変換がLayerNormalizedLSTMCellSimpleに対して記述されます。

  8. TensorFlow Lite 変換インフラストラクチャの残りの部分 (定義されたすべてのMLIR パスや TensorFlow Lite フラットバッファへの最終エクスポートを含む) は再利用できます。

既知の問題/制限事項

  1. 現在、ステートレス Keras LSTM (Keras のデフォルト動作) の変換のみがサポートされています。ステートフル Keras LSTM への変換は今後の課題です。
  2. 基礎となるステートレス Keras LSTM レイヤーを使用してステートフル Keras LSTM レイヤーをモデル化し、ユーザー プログラムで明示的に状態を管理することは依然として可能です。このような TensorFlow プログラムは、ここで説明する機能を使用して TensorFlow Lite に変換できます。
  3. 双方向 LSTM は現在、TensorFlow Lite の 2 つの UnidirectionalSequenceLSTM 操作としてモデル化されています。これは、単一の BidirectionSequenceLSTM 演算に置き換えられます。