概要
このページでは、複合オペレーションを変換するために必要な設計と手順について説明します。 LiteRT での融合演算について説明します。このインフラストラクチャは、 TensorFlow での任意の複合演算の変換をサポート 対応する融合演算（LiteRT）にマッピングします。
このインフラストラクチャの使用例として、TensorFlow RNN 演算の融合、 LiteRT。詳しくはこちらをご覧ください。
融合オペレーションとは
TensorFlow の演算は、 tf.add することも、 他のプリミティブ操作から構成されます。 tf.einsum。プリミティブ 演算は TensorFlow グラフに 1 つのノードとして表示され、一方で複合演算は 演算は TensorFlow グラフ内のノードの集まりです。実行中の 複合オペレーションは、その構成要素である各プリミティブを実行するのと同等です。 必要があります。
融合オペレーションは、単一のオペレーションに それぞれのプリミティブ演算によって実行される 複合オペレーションがあります。
融合オペレーションのメリット
融合されたオペレーションは基盤となるカーネルのパフォーマンスを最大化するために存在します。 全体的な計算を最適化し、メモリを削減することで、実装を あります。これは、特に低レイテンシの推論ワークロードに非常に有益です。 リソースの制約があるモバイルプラットフォームに 対応できます
融合されたオペレーションは、複雑なリソースを定義したり、 量子化などの高度な変換は不可能であり、他の方法では 細かいレベルで行うのは困難です
LiteRT には、次の理由により、融合されたオペレーションのインスタンスが多数あります。 必要がありますこれらの融合されたオペレーションは通常、 数個の演算を実行します。Google Cloud Storage の複合オペレーション LiteRT で単一の融合演算として実装される TensorFlow 単方向および双方向シーケンスなどのさまざまな RNN 演算を LSTM、畳み込み（conv2d、bias add、relu）、全接続（matmul、bias add、 Relu など）が含まれます。LiteRT では、現在のところ LSTM 量子化は 実装されています。
オペレーションの融合の課題
複合演算を TensorFlow から融合演算に変換するには、 LiteRT は難しい問題です。理由は次のとおりです。
複合演算は TensorFlow グラフで一連の 境界が明確に定義されていないプリミティブ操作です。非常に （パターン マッチングなどによって）サブグラフを特定するのは困難です。 対応する複合オペレーションです
融合されたものをターゲットとする TensorFlow 実装が複数存在する場合があります。 LiteRT オペレーション。たとえば、LSTM の実装は多数あります。 （Keras、Babelfish/lingvo など）で記述されており、それぞれが すべて変換することもできますが、それらはすべて LiteRT と同じ融合 LSTM 演算です。
そのため、融合されたオペレーションの変換は非常に困難でした。
複合演算から TFLite カスタム オペレーションへの変換（推奨）
複合オペレーションを tf.function でラップする
多くの場合、モデルの一部を単一のオペレーションにマッピングできます。
TFLite.これにより、最適化された実装を記述する際のパフォーマンスが向上します。
使用できます。TFLite で融合オペレーションを作成できるようにするには、
グラフの中で融合演算を表す部分を特定し
tf.function を、
"experimental_ implementss"属性を持つ tf.function に、属性を
値 tfl_fusable_op と値 true。カスタム オペレーションで
同じ「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))
コンバータで allow_custom_ops を設定する必要はありません。
tfl_fusable_op 属性は、すでにこれに該当します。
カスタム演算を実装して TFLite インタープリタに登録する
融合オペレーションを TFLite カスタム オペレーションとして実装します。以下をご覧ください。 こちらの手順をご覧ください。
演算を登録する際の名前は、その名前と類似している必要があります。
実装署名の 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, ®);
複合オペレーションから融合オペレーションへの変換（Advanced）
TensorFlow の複合オペレーションを LiteRT の融合オペレーションは次のとおりです。
複合オペレーションを tf.function でラップする
TensorFlow モデルのソースコードで、複合オブジェクトを特定して抽象化する すべてのオペレーションを tf.function を、 experimental_implements 関数アノテーションを使用します。エンベディング ルックアップの例をご覧ください。「 関数でインターフェースを定義し、その引数を使用して 変換ロジックを渡します。
コンバージョン コードを記述する
変換コードは、関数のインターフェースに従って
implements アノテーション。埋め込みの融合の例を参照する。
検索します。概念的には、変換コードは
このインターフェースの実装例を示します。
prepare-composite-functions のパスで、コンバージョン 提供します。
より高度な使い方では、インフラストラクチャの複雑な変換を実装して、 複合演算のオペランドを渡して、融合された あります。詳細については、Keras LSTM. ご覧ください。
LiteRT に変換する
こちらの TFLiteConverter.from_saved_model LiteRT に変換する API。
詳細
ここでは、融合への変換における全体的な設計の概要を説明します。 LiteRT での操作を行えます。
TensorFlow での演算の作成
tf.function の使用 新しい experimental_implements 使用することで、ユーザーは新しいオペレーションを TensorFlow のプリミティブ演算を実行し、 実装します。次の機能を備えているため、非常に便利です。
- 基盤となる複合オペレーションの境界が明確に定義された TensorFlow のグラフ。
- このオペレーションが実装するインターフェースを明示的に指定します。「 引数として tf.function 各引数は、このインターフェースの引数に対応しています。
例として、次の要素を実装するために定義された複合オペレーションについて考えてみましょう。 エンベディング ルックアップです。これは LiteRT の融合オペレーションにマッピングされます。
@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 を 上の図のように、インフラストラクチャを このようなオペレーションを特定して変換し、LiteRT オペレーションを融合します。
LiteRT コンバータの拡張
今年初めにリリースされた LiteRT コンバータは、 TensorFlow モデルをグラフとしてインポートする。このグラフでは、すべての変数がそれぞれのモデルに 対応する定数値を使用します。この操作はオペレーションの融合では機能しません。 すべての関数がインライン化されているので、変数を 使用します。
Google Cloud の
tf.function を、
変換処理中に experimental_implements 特徴が使用された場合、
変換プロセスの後半まで保持する必要があります。
そのため、TensorFlow をインポートして変換する新しいワークフローを実装しました。 コンバータ内でモデルを作成して、複合演算の融合のユースケースをサポートします。 具体的には、次のような新機能が追加されました。
- TensorFlow の保存済みモデルを MLIR にインポートする
- 複合オペレーションを融合する
- 変数の可変性の分析
- すべての読み取り専用変数を固定する
これにより、入力値を表す関数を使用して、オペレーションの融合を実行できます。 複合オペレーションを実行しています。
Operation Fusion の実装
オペレーション フュージョン パスをさらに詳しく見てみましょう。このパスは、 次のとおりです。
- MLIR モジュール内のすべての関数をループする。
- 関数に tf._Implements 属性がある場合は、 適切なオペレーション融合ユーティリティを呼び出します。
- Fusion ユーティリティは、関数のオペランドと （変換のインターフェースとして機能する）とともに、 関数本体と同等の関数本体を 説明します。
- 多くの場合、置き換えられる本文には 説明します。これらはモデルのいくつかの静的変換に対応しています。 融合演算のオペランドを取得するために、関数のオペランドを渡しています。 これらの計算はすべて一定に折りたたむことができるため、 エクスポートされたフラットバッファに格納されていて、 あります。
メイン ワークフローを示すパスのコード スニペットを以下に示します。
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 */
}
この複合オペレーションを融合オペレーションにマッピングするコード スニペットを以下に示します。 関数を変換インターフェースとして使用して、LiteRT でオペレーションを実行します。
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());
}