TensorFlow 操作の融合

概要

このページでは、TensorFlow の複合演算を TensorFlow Lite の融合演算に変換するために必要な設計と手順について説明します。このインフラストラクチャは汎用であり、TensorFlow の複合演算から TensorFlow Lite の対応する融合演算への変換をサポートします。

このインフラストラクチャの使用例は、ここで詳しく説明されているように、TensorFlow Lite への TensorFlow RNN 操作の融合です。

融合操作とは

描画

TensorFlow オペレーションは、プリミティブ オペレーション (例: tf.add)にすることも、他のプリミティブ オペレーション (例: tf.einsum )から構成することもできます。プリミティブ操作は TensorFlow グラフ内の単一ノードとして表示されますが、複合操作は TensorFlow グラフ内のノードのコレクションです。複合操作の実行は、その構成要素である各プリミティブ操作を実行することと同等です。

融合演算は、対応する複合演算内の各基本演算によって実行されるすべての計算を包含する単一の演算に対応します。

融合操作の利点

融合オペレーションは、全体的な計算を最適化し、メモリ フットプリントを削減することにより、基礎となるカーネル実装のパフォーマンスを最大化するために存在します。これは、特に低レイテンシの推論ワークロードやリソースに制約のあるモバイル プラットフォームにとって非常に価値があります。

融合操作は、量子化などの複雑な変換を定義するための高レベルのインターフェイスも提供します。これがなければ、より詳細なレベルで実行できないか、非常に困難です。

TensorFlow Lite には、上で述べた理由により、融合された操作のインスタンスが多数あります。これらの融合操作は通常、ソース TensorFlow プログラムの複合操作に対応します。 TensorFlow Lite で単一の融合演算として実装される TensorFlow の複合演算の例には、単方向および双方向シーケンス LSTM、畳み込み (conv2d、bias add、relu)、完全接続 (matmul、bias add、relu) などのさまざまな RNN 演算が含まれます。 。 TensorFlow Lite では、LSTM 量子化は現在、融合された LSTM 操作でのみ実装されています。

融合オペレーションの課題

TensorFlow から TensorFlow Lite の融合演算に複合演算を変換するのは難しい問題です。それの訳は：

  1. 複合操作は、TensorFlow グラフでは、明確に定義された境界のない一連のプリミティブ操作として表されます。このような複合演算に対応するサブグラフを (パターン マッチングなどにより) 識別することは非常に困難な場合があります。

  2. 融合された TensorFlow Lite 操作をターゲットとする TensorFlow 実装が複数存在する場合があります。たとえば、TensorFlow には多くの LSTM 実装 (Keras、Babelfish/lingvo など) があり、これらのそれぞれは異なるプリミティブ操作で構成されていますが、それらはすべて TensorFlow Lite の同じ融合された LSTM 操作に変換できます。

そのため、融合された操作の変換は非常に困難であることが判明しています。

複合操作をtf.functionでラップします。

多くの場合、モデルの一部を TFLite の単一の操作にマップできます。これは、特定の操作に最適化された実装を作成する際のパフォーマンスの向上に役立ちます。 TFLite で融合操作を作成できるようにするには、融合操作tf.function tf.functionラップし、値trueの属性値tfl_fusable_opを持つ tf.function にラップします。カスタム操作が属性を取る場合は、それらを同じ「experimental_implements」の一部として渡します。

例、

def get_implements_signature():
  implements_signature = [
    # 'name' will be used as a name for the operation.
    'name: "my_custom_fused_op"',
    # attr "tfl_fusable_op" is required to be set with true value.
    'attr {key: "tfl_fusable_op" value { b: true } }',
    # Example attribute "example_option" that the op accepts.
    'attr {key: "example_option" value { i: %d } }' % 10
  ]
  return ' '.join(implements_signature)

@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
  # An empty function that represents pre/post processing example that
  # is not represented as part of the Tensorflow graph.
  output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
  output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
  return output_1, output_2

class TestModel(tf.Module):
  def __init__(self):
    super(TestModel, self).__init__()
    self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
    self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
  ])
  def simple_eval(self, input_a, input_b):
    return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))

tfl_fusable_op属性がすでにこれを暗示しているため、コンバーターでallow_custom_ops設定する必要がないことに注意してください。

カスタム演算を実装し、TFLite インタープリターに登録する

融合操作を TFLite カスタム操作として実装します。手順を参照してください。

op を登録する名前は、implements シグネチャのname属性で指定された名前と似ている必要があることに注意してください。

この例の操作の例は次のとおりです。

  TfLiteRegistration reg = {};
  // This name must match the name specified in the implements signature.
  static constexpr char kOpName[] = "my_custom_fused_op";
  reg.custom_name = kOpName;
  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

複合操作から融合操作への変換 (上級)

TensorFlow 複合操作を TensorFlow Lite 融合操作に変換するための全体的なアーキテクチャは次のとおりです。

描画

複合操作をtf.functionでラップします。

TensorFlow モデルのソース コードで、複合操作を特定し、 experimental_implements関数アノテーションを使用してtf.functionに抽象化します。埋め込みルックアップの例を参照してください。この関数はインターフェイスを定義し、その引数は変換ロジックの実装に使用する必要があります。

変換コードを書く

変換コードは、 implementsアノテーションを使用して関数のインターフェイスごとに記述されます。ルックアップを埋め込むための融合例を参照してください。概念的には、変換コードはこのインターフェイスの複合実装を融合された実装に置き換えます。

prepare-composite-functions パスで、変換コードをプラグインします。

より高度な使用法では、融合演算のオペランドを導出するために、複合演算のオペランドの複雑な変換を実装することが可能です。 Keras LSTMを参照してください。例として変換コードを示します。

TensorFlow Lite に変換する

TFLiteConverter.from_saved_model API を使用して TensorFlow Lite に変換します。

フードの下

次に、TensorFlow Lite での融合操作への変換における全体的な設計の高レベルの詳細について説明します。

TensorFlow での操作の合成

tf.functionをExperimental_implements関数属性とともに使用すると、ユーザーは TensorFlow プリミティブ操作を使用して新しい操作を明示的に作成し、結果の複合操作が実装するインターフェイスを指定できます。これは以下を提供するので非常に便利です。

  1. 基礎となる TensorFlow グラフ内の複合操作の明確に定義された境界。
  2. この操作が実装するインターフェイスを明示的に指定します。 tf.functionの引数は、このインターフェイスの引数に対応します。

例として、埋め込み検索を実装するために定義された複合操作を考えてみましょう。これは、TensorFlow Lite の融合操作にマッピングされます。

  @tf.function(
        experimental_implements="embedding_lookup")
    def EmbFprop(embs, ids_vec):
      """Embedding forward prop.

      Effectively, it computes:
        num = size of ids_vec
        rets = zeros([num, embedding dim])
        for i in range(num):
          rets[i, :] = embs[ids_vec[i], :]
        return rets

      Args:
        embs: The embedding matrix.
        ids_vec: A vector of int32 embedding ids.

      Returns:
        The result of embedding lookups. A matrix of shape
        [num ids in ids_vec, embedding dims].
      """
      num = tf.shape(ids_vec)[0]
      rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))

      def EmbFpropLoop(i, embs, ids_vec, rets):
        # row_id = ids_vec[i]
        row_id = tf.gather(ids_vec, i)
        # row = embs[row_id]
        row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
        # rets[i] = row
        rets = inplace_ops.alias_inplace_update(rets, [i], row)
        return embs, ids_vec, rets

      _, _, rets = functional_ops.For(
          start=0,
          limit=num,
          delta=1,
          inputs=[embs, ids_vec, rets],
          body=EmbFpropLoop,
          rewrite_with_while=compiled)
      if len(weight_shape) > 2:
        rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
      return rets

上に示したように、 tf.functionを介してモデルで複合演算を使用するようにすることで、そのような演算を識別し、融合された TensorFlow Lite 演算に変換するための一般的なインフラストラクチャを構築することが可能になります。

TensorFlow Lite コンバータの拡張

今年初めにリリースされた TensorFlow Lite コンバータは、すべての変数が対応する定数値に置き換えられたグラフとして TensorFlow モデルをインポートすることのみをサポートしていました。このようなグラフでは変数を定数に変えることができるようにすべての関数がインライン化されているため、これは演算融合では機能しません。

変換プロセス中に、 experimental_implements機能を備えたtf.functionを利用するには、変換プロセスの後半まで関数を保存する必要があります。

そのため、複合操作フュージョンのユースケースをサポートするために、コンバーターに TensorFlow モデルをインポートして変換する新しいワークフローを実装しました。具体的には、追加された新機能は次のとおりです。

  1. TensorFlow で保存されたモデルを MLIR にインポートする
  2. ヒューズ複合操作
  3. 変数の可変性分析
  4. すべての読み取り専用変数を凍結する

これにより、関数のインライン化と変数の凍結の前に、複合演算を表す関数を使用して演算融合を実行できるようになります。

オペレーションフュージョンの実装

オペレーション フュージョン パスをさらに詳しく見てみましょう。このパスは次のことを行います。

  1. MLIR モジュール内のすべての関数をループします。
  2. 関数に tf._implements 属性がある場合、その属性値に基づいて、適切な操作融合ユーティリティが呼び出されます。
  3. 操作融合ユーティリティは、関数のオペランドと属性 (変換のインターフェイスとして機能します) を操作し、関数の本体を、融合された操作を含む同等の関数本体に置き換えます。
  4. 多くの場合、置き換えられたボディには、融合された操作以外の操作が含まれます。これらは、融合演算のオペランドを取得するための、関数のオペランドの静的変換に対応します。これらの計算はすべて定数で折り畳むことができるため、融合された操作のみが存在するエクスポートされたフラットバッファには存在しません。

以下は、メインのワークフローを示すパスのコード スニペットです。

void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
                                                        StringAttr attr) {
  if (attr.getValue() == "embedding_lookup") {
    func.eraseBody();
    func.addEntryBlock();
    // Convert the composite embedding_lookup function body to a
    // TFLite fused embedding_lookup op.
    ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
    if (failed(convert_embedded_lookup.VerifySignature())) {
      return signalPassFailure();
    }
    convert_embedded_lookup.RewriteFunc();
  } else if (attr.getValue() == mlir::TFL::kKerasLstm) {
     func.eraseBody();
     func.addEntryBlock();
     OpBuilder builder(func.getBody());
     if (failed(ConvertKerasLSTMLayer(func, &builder))) {
       return signalPassFailure();
     }
  } else if (.....) /* Other fusions can plug in here */
}

以下は、関数を変換インターフェイスとして利用して、この複合演算を TensorFlow Lite の融合演算にマッピングするコード スニペットです。

void RewriteFunc() {
    Value lookup = func_.getArgument(1);
    Value value = func_.getArgument(0);
    auto output_type = func_.getType().getResult(0);

    OpBuilder builder(func_.getBody());
    auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
        func_.getLoc(), output_type, lookup, value);

    builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
  }