TensorFlow Lite のビルトイン演算子ライブラリがサポートする TensorFlow 演算子は制限されているため、すべてのモデルは変換可能ではありません。詳細は、演算子の互換性をご覧ください。
変換を可能にするためには、カスタム演算子を使用します。カスタム演算子とは、TensorFlow Lite でサポートされていない TensorFlow 演算子のカスタム実装です。また、サポートされていない（またはサポートされている）一連の TensorFlow 演算子を、融合して最適化された単一のカスタム演算子として使用できます。詳細は演算子の融合を参照してください。
カスタム演算子を使用するには、次の 4 つのステップに従います。
TensorFlow モデルを作成する。Saved Model（または Graph Def）が正しく名付けられた TensorFlow Lite 演算子を参照するようにします。
TensorFlow Lite モデルに変換する。モデルを正しく変換するには、適切な TensorFlow Lite コンバータ属性を指定します。
演算子を作成して登録する。TensorFlow Lite ランタイムに、グラフの演算子とパラメータを実行可能な C/C++ コードにマッピングする方法を認識させるために行います。
演算子をテストしプロファイリングする。カスタム演算子のみをテストする場合は、カスタム演算子のみでモデルを作成し、benchmark_model プログラムを使用することをお勧めします。
では、カスタム演算子 tf.atan（Atan と名付けられています。TensorFlow モデルを作成するのセクションをご覧ください）を使ってモデルを実行するエンドツーエンドの例を見てみましょう。このカスタム演算子は TensorFlow ではサポートされていますが、TensorFlow Lite ではサポートされていません。
注意: tf.atan 関数はカスタム演算子ではありません。TensorFlow と TensorFlow Lite の両方でサポートされている通常の演算子ですが、単純なワークフローを実演するために、次の例ではカスタム演算子であると仮定しています。
TensorFlow Text 演算子はカスタム演算子の例です。コードサンプルについては、TF Text を TF Lite に変換チュートリアルをご覧ください。
例: カスタム Atan 演算子
TensorFlow Lite にはない TensorFlow 演算子をサポートする例を見てみましょう。Atan 演算子を使用し、関数 y = atan(x + offset) の非常に簡単なモデルを構築するとします。この関数の offset はトレーニング可能です。
TensorFlow モデルを作成する
次は、単純な TensorFlow モデルをトレーニングするコードスニペットです。このモデルには、Atan というカスタム演算子のみが含まれています。このカスタム演算は関数 y = atan(x + offset) で、offset はトレーニング可能です。
import tensorflow as tf
# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-1.4288993, 0.98279375, 1.2490457, 1.2679114, 1.5658458]
offset = tf.Variable(0.0)
# Define a simple model which just contains a custom operator named `Atan`
@tf.function(input_signature=[tf.TensorSpec.from_tensor(tf.constant(x))])
def atan(x):
return tf.atan(x + offset, name="Atan")
# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
with tf.GradientTape() as t:
predicted_y = atan(x)
loss = tf.reduce_sum(tf.square(predicted_y - y))
grads = t.gradient(loss, [offset])
optimizer.apply_gradients(zip(grads, [offset]))
for i in range(1000):
train(x, y)
print("The actual offset is: 1.0")
print("The predicted offset is:", offset.numpy())
The actual offset is: 1.0
The predicted offset is: 0.99999905
この時点では、デフォルトのコンバータフラグを使って TensorFlow Lite モデルを生成しようとすると、次のようなエラーメッセージが発生します。
Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.
TensorFlow Lite モデルに変換する
カスタム演算子を使った TensorFlow Lite モデルを作成しましょう。次に示すように、コンバータの属性をallow_custom_opsに設定してください。
converter = tf.lite.TFLiteConverter.from_concrete_functions([atan.get_concrete_function()], atan) <b>converter.allow_custom_ops = True</b> tflite_model = converter.convert()
この時点では、以下のようなコマンドを使ってデフォルトのインタプリタで実行しようとすると:
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
以下のようなエラーメッセージが発生します。
Encountered unresolved custom op: Atan.
演算を作成して登録する
すべての TensorFlow Lite 演算子（カスタムとビルトインの両方）は、次の 4 つの関数で構成される単純な Pure C インターフェースを使って定義されています。
typedef struct {
void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
void (*free)(TfLiteContext* context, void* buffer);
TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
} TfLiteRegistration;
TfLiteContext と TfLiteNode の詳細については、common.h をご覧ください。前者はエラー報告ファシリティと、すべてのテンソルを含むグローバルオブジェクトへのアクセスを提供し、後者は実装が入力と出力にアクセスできるようにします。
インタプリタがモデルを読み込む際、グラフの各ノード当たりinit()が一度呼び出されます。つまりグラフで演算が 2 回以上使用される場合、init()は 2 回以上呼び出されることになります。カスタム演算子の場合、パラメータ名を値にマッピングする flexbuffer が含まれる構成バッファが提供されます。このバッファは、インタプリタが演算子のパラメータをパースするため、組み込み演算子の場合は空です。ステートが必要なカーネル実装はここで初期化して、オーナーシップを呼び出し元に移譲します。各 init()呼び出しにはfree()への対応する呼び出しがあるため、実装は、init()に割り当てられていた可能性のあるバッファを利用できるようになります。
入力テンソルのサイズが変更されるたび、インタプリタはグラフを通過して実装に変更を通知します。このため、実装は、内部バッファのサイズを変更し、入力形状と型の妥当性を検証し、出力形状を再計算するチャンスを得ることができます。このプロセスはすべてprepare()を介して行われるため、実装はnode->user_dataを使用してステートにアクセスすることができます。
最後に、推論が実行されるたび、インタプリタは、invoke()を呼び出してグラフをトラバースし、ここでもステートをnode->user_dataとして使用することができます。
カスタム演算子は、定義済みの 4 つの関数と、通常は次のようなグローバル登録関数によって、組み込み演算子とまったく同じように実装できます。
namespace tflite {
namespace ops {
namespace custom {
TfLiteRegistration* Register_MY_CUSTOM_OP() {
static TfLiteRegistration r = {my_custom_op::Init,
my_custom_op::Free,
my_custom_op::Prepare,
my_custom_op::Eval};
return &r;
}
} // namespace custom
} // namespace ops
} // namespace tflite
登録は自動ではなく、明示的なRegister_MY_CUSTOM_OP呼び出しを行う必要があることに注意してください。標準のBuiltinOpResolver（:builtin_opsターゲットから利用可能）が組み込み関数を処理する間に、別のカスタムライブラリでカスタム演算子を収集する必要があります。
TensorFlow Lite ランタイムでカーネルを定義する
TensorFlow Lite でこの演算子を使用できるようにするには、2 つの関数（PrepareおよびEval）を定義し、TfLiteRegistrationを構築する必要があります。
TfLiteStatus AtanPrepare(TfLiteContext* context, TfLiteNode* node) {
using namespace tflite;
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input = GetInput(context, node, 0);
TfLiteTensor* output = GetOutput(context, node, 0);
int num_dims = NumDimensions(input);
TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
for (int i=0; i<num_dims; ++i) {
output_size->data[i] = input->dims->data[i];
}
return context->ResizeTensor(context, output, output_size);
}
TfLiteStatus AtanEval(TfLiteContext* context, TfLiteNode* node) {
using namespace tflite;
const TfLiteTensor* input = GetInput(context, node, 0);
TfLiteTensor* output = GetOutput(context, node, 0);
float* input_data = GetTensorData<float>(input);
float* output_data = GetTensorData<float>(output);
size_t count = 1;
int num_dims = NumDimensions(input);
for (int i = 0; i < num_dims; ++i) {
count *= input->dims->data[i];
}
for (size_t i=0; i<count; ++i) {
output_data[i] = atan(input_data[i]);
}
return kTfLiteOk;
}
TfLiteRegistration* Register_ATAN() {
static TfLiteRegistration r = {nullptr, nullptr, AtanPrepare, AtanEval};
return &r;
}
OpResolver を初期化する際に、カスタム演算子をレゾルバに追加します（以下の例をご覧ください）。これにより、演算子は TensorFlow Lite に登録され、TensorFlow Lite で新しい実装として使用できるようになります。TfLiteRegistration の最後の 2 つの引数は、カスタム演算子に定義した AtanPrepare と AtanEval 関数に対応しています。AtanInit と AtanFree 関数を使用して、それぞれ演算子に使用される変数を初期化して容量を解放すると、TfLiteRegistration の最初の 2 つの引数に追加されます。この例ではこれらの引数は nullptr に設定されます。
カーネルライブラリで演算子を登録する
では、カーネルライブラリを使って演算子を登録する必要があります。これは、OpResolverを使って行います。舞台裏では、インタプリタは、モデルの各演算子を実行するために割り当てられるカーネルのライブラリを読み込みます。デフォルトのライブラリには組み込みカーネルのみが含まれますが、これをカスタムライブラリに置き換えたり拡張することが可能です。
演算子のコードと名前を実際のコードに変換するOpResolverクラスは、次のように定義されます。
class OpResolver {
virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
virtual TfLiteRegistration* FindOp(const char* op) const = 0;
virtual void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration) = 0;
virtual void AddCustom(const char* op, TfLiteRegistration* registration) = 0;
};
通常の使用では、BuiltinOpResolverを使用して次のように書く必要があります。
tflite::ops::builtin::BuiltinOpResolver resolver;
上記で作成したカスタム演算子を追加するには、（レゾルバをInterpreterBuilderに渡す前に）AddOpを呼び出します。
resolver.AddCustom("Atan", Register_ATAN());
組み込み演算子のセットが大きすぎる場合、演算子の特定のサブセットに基づいて、新しいOpResolverをコード生成することができます。おそらく特定のモデルに含まれる演算子のみが含まれます。これが、TensorFlow の選択的登録に相当するものです（また、この単純なバージョンは、toolsディレクトリに提供されています）。
カスタム演算を Java で定義する場合、現時点では、独自のカスタム JNI レイヤーを構築し、この jni コードに独自の AAR をコンパイルする必要があります。同様に、Python でこれらの演算を定義する場合、Python ラッパーコードに登録を配置することができます。
上記に類似するプロセスは、単一の演算子ではなく一連の演算子をサポートするために使用することができます。必要な数のAddCustom演算子を追加してください。また、BuiltinOpResolverを使用することで、AddBuiltinを使用して、組み込みの実装をオーバーライドできます。
演算子をテストしてプロファイリングする
TensorFlow Lite のベンチマークツールで演算子をプロファイリングするには、TensorFlow Lite 用のベンチマークモデルツールを使用できます。テスト目的により、該当する AddCustom 呼び出しを（上記に示すとおり）register.cc に追加すれば、TensorFlow Lite のローカルビルドにカスタム演算子を認識させることができます。
ベストプラクティス
メモリの割り当てと解除の最適化には注意してください。
Prepareにメモリを割り当てる方が、Invokeに割り当てるよりも効果的であり、ループの前に割り当てる方が、各イテレーションで割り当てるよりも有効です。自分で割り当てるよりも一時テンソルデータを使用してください（項目 2 参照）。できるだけコピーよりも、ポインタ/参照を使用してください。演算を通じてデータ構造が永続する場合、一時テンソルを使用してメモリを事前に割り当てておくことをお勧めします。ほかの関数でテンソルインデックスを参照するには、OpData 構造を使用する必要がある場合があります。畳み込み用のカーネルの例をご覧ください。サンプルのコードスニペットを次に示します。
auto* op_data = reinterpret_cast<OpData*>(node->user_data); TfLiteIntArrayFree(node->temporaries); node->temporaries = TfLiteIntArrayCreate(1); node->temporaries->data[0] = op_data->temp_tensor_index; TfLiteTensor* temp_tensor = &context->tensors[op_data->temp_tensor_index]; temp_tensor->type = kTfLiteFloat32; temp_tensor->allocation_type = kTfLiteArenaRw;メモリの浪費がかさばらないようであれば、実行のイテレーションごとに動的に割り当てられる
std::vectorよりも、静的な固定サイズの配列（またはResizeの事前割り当て済みのstd::vector）を使用することをお勧めします。存在しない標準ライブラリコンテナテンプレートをインスタンス化しないようにしてください。バイナリサイズが大きくなります。たとえば、ほかのカーネルに存在しない
std::mapが演算に必要な場合、直接インデックスマッピングでstd::vectorを使用すれば、バイナリサイズを小さくしたまま機能させることが可能です。ほかのカーネルが何を使用して洞察を得ているのかご覧ください。mallocが返すメモリへのポインタを確認してください。このポインタがnullptrである場合、そのポイントを使って演算を実行してはいけません。関数でmallocを行い、エラーが存在する場合は、終了する前にメモリを解放してください。特定の条件を確認するには、
TF_LITE_ENSURE(context, condition)を使います。TF_LITE_ENSUREが使用されるときには、コードはメモリを残してはいけません。これらのマクロは、リークするリソースが割り当てられる前に使用される必要があります。