Google AI Edge Portal のご紹介: エッジ AI を大規模にベンチマークします。限定公開プレビュー中にアクセスをリクエストするには、登録してください。

このページは Cloud Translation API によって翻訳されました。

カスタム演算子

LiteRT 組み込み演算子ライブラリは、限られた数の TensorFlow 演算子しかサポートしていないため、すべてのモデルを変換できるわけではありません。詳細については、演算子の互換性をご覧ください。

変換を許可するために、ユーザーは LiteRT でサポートされていない TensorFlow オペレーターの独自のカスタム実装（カスタムオペレーター）を提供できます。サポートされていない（またはサポートされている）一連の TensorFlow 演算子を 1 つの融合された最適化されたカスタム演算子に結合する場合は、演算子の融合を参照してください。

カスタムオペレータの使用は、次の 4 つの手順で構成されます。

TensorFlow モデルを作成します。Saved Model（または Graph Def）が、正しい名前の LiteRT オペレーターを参照していることを確認します。
LiteRT モデルに変換します。モデルを正常に変換するには、適切な LiteRT コンバータ属性を設定してください。
オペレーターを作成して登録します。これは、LiteRT ランタイムがグラフ内の演算子とパラメータを実行可能な C/C++ コードにマッピングする方法を認識できるようにするためです。
オペレーターをテストしてプロファイリングします。カスタムオペレーターのみをテストする場合は、カスタムオペレーターのみを含むモデルを作成し、benchmark_model プログラムを使用することをおすすめします。

TensorFlow でサポートされていて、LiteRT ではサポートされていないカスタム演算子 tf.atan（Atan という名前。TensorFlow モデルを作成するを参照）を使用してモデルを実行するエンドツーエンドの例を見てみましょう。

TensorFlow Text 演算子は、カスタム演算子の例です。コード例については、TF Text を LiteRT に変換するチュートリアルをご覧ください。

例: カスタム `Atan` 演算子

LiteRT にない 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

この時点で、デフォルトのコンバータフラグを使用して LiteRT モデルを生成しようとすると、次のエラーメッセージが表示されます。

Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.

LiteRT モデルに変換する

次の例に示すように、コンバータ属性 allow_custom_ops を設定して、カスタムオペレータを含む LiteRT モデルを作成します。

converter = tf.lite.TFLiteConverter.from_concrete_functions([atan.get_concrete_function()], atan)
converter.allow_custom_ops = True
tflite_model = converter.convert()

この時点で、次のようなコマンドを使用してデフォルトのインタープリタで実行すると、

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

エラーは引き続き表示されます。

Encountered unresolved custom op: Atan.

オペレーターを作成して登録します。

#include "third_party/tensorflow/lite/c/c_api.h"
#include "third_party/tensorflow/lite/c/c_api_opaque.h"

LiteRT カスタムオペレーターは、不透明型（TfLiteOperator）と関連関数で構成されるシンプルな純粋な C API を使用して定義されます。

TfLiteOperator は不透明型です。

typedef struct TfLiteOperator TfLiteOperator;

TfLiteOperator には、オペレーターの ID と実装が保存されます。（演算子は、演算子を呼び出すノードの LiteRT グラフノードに保存されるオペランドとは異なります）。

この型のインスタンスは TfLiteOperatorCreate の呼び出しで構築され、TfLiteOperatorDelete の呼び出しで破棄できます。

オペレータの ID は、コンストラクタ関数 TfLiteOperatorCreate のパラメータで設定されます。

TfLiteOperator*
TfLiteOperatorCreate(
    TfLiteBuiltinOperator builtin_code,  // Normally `TfLiteBuiltinCustom`.
    const char* custom_name,  // The name of the custom op.
    int version  // Normally `1` for the first version of a custom op.
);

オペレーター実装では、次のシグネチャで「メソッド」を定義できます。これらのメソッドはすべて省略可能ですが、オペレーターが正常に評価されるためには、オペレーターの実装で少なくとも Prepare メソッドと Invoke メソッドを定義して設定（セッター関数を使用）する必要があります。

// Initializes the op from serialized data.
void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length);

// Deallocates the op.
// The pointer `buffer` is the data previously returned by an Init invocation.
void Free(TfLiteOpaqueContext* context, void* buffer);

// Called when the inputs that this node depends on have been resized.
TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);

// Called when the node is executed. (Should read node inputs and write to
// node outputs).
TfLiteStatus Invoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);

// Retrieves the async kernel.
TfLiteAsyncKernel AsyncKernel(TfLiteOpaqueContext* context,
                              TfLiteOpaqueNode* node);

TF Lite カスタムオペレーション API はアドレスのみを使用するため、オペレーション実装の関数名（または C++ の名前空間接頭辞）は、上記のコードスニペットの関数名と一致する必要はありません。実際、匿名 Namespace で宣言するか、静的関数として宣言することをおすすめします。

ただし、これらの関数名に名前空間または接頭辞として演算子名を含めることをおすすめします。

C++

namespace my_namespace::my_custom_op {
  void* Init(TfLiteOpaqueContext* context,
             const char* buffer, size_t length) { ... }
  // ... plus definitions of Free, Prepare, and Invoke ...
}

C

void* MyCustomOpInit(TfLiteOpaqueContext* context,
                     const char* buffer, size_t length) { ... }
// ... plus definitions of MyCustomOpFree, MyCustomOpPrepare, and
// MyCustomOpInvoke.

これは C API であるため、これらの「メソッド」は TfLiteOperator 型の C 関数ポインタとして実装されます。これは、実装関数のアドレスを対応するセッター関数 TfLiteOperatorSetMethodName に渡すことで設定されます。

void TfLiteOperatorSetInit(
    TfLiteOperator* operator,
    void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
                  size_t length));
void TfLiteOperatorSetFree(
    TfLiteOperator* operator,
    void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteOperatorSetPrepare(
    TfLiteOperator* operator,
    TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
                            TfLiteOpaqueNode* node));
void TfLiteOperatorSetInvoke(
    TfLiteOperator* operator,
    TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
                           TfLiteOpaqueNode* node));
void TfLiteOperatorSetAsyncKernel(
    TfLiteOperator* operator,
    struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
                                              TfLiteOpaqueNode* node));

TfLiteContext と TfLiteNode の詳細については、common.h をご覧ください。TfLiteContext は、エラー報告機能と、すべてのテンソルを含むグローバルオブジェクトへのアクセスを提供します。TfLiteNode を使用すると、オペレーター実装で入出力にアクセスできます。

インタープリタがモデルを読み込むと、グラフ内の各ノードに対して Init() メソッドを 1 回呼び出します。グラフ内で op が複数回使用されている場合、特定の Init() が複数回呼び出されます。カスタムオペレーションの場合、パラメータ名とその値をマッピングする flexbuffer を含む構成バッファが提供されます。インタープリタがすでに op パラメータを解析しているため、組み込み op のバッファは空です。状態を必要とするカーネル実装は、ここで状態を初期化し、所有権を呼び出し元に転送する必要があります。Init() 呼び出しごとに、対応する Free() 呼び出しが行われ、実装で Init() で割り当てたバッファを破棄できるようになります。

入力テンソルがリサイズされるたびに、インタープリタはグラフを通過して、変更の実装に通知します。これにより、内部バッファのサイズ変更、入力形状と型の有効性の確認、出力形状の再計算が可能になります。これはすべて Prepare() メソッドで行われ、実装では TfLiteOpaqueNodeGetUserData(node) を使用して状態にアクセスできます。

最後に、推論が実行されるたびに、インタープリタは Invoke() メソッドを呼び出すグラフをトラバースします。ここでも、状態は TfLiteOpaqueNodeGetUserData(node) として使用できます。

カスタムオペレーションは、これらの「メソッド」関数を定義し、TfLiteOperatorCreate と関連するセッターメソッドを呼び出して構築された TfLiteOperator のインスタンスを返す関数を定義することで実装できます。

C++

namespace my_namespace::my_custom_op {
  namespace {
    void* Init(TfLiteOpaqueContext* context,
               const char* buffer, size_t length) { ... }
    void Free(TfLiteOpaqueContext* context, void* buffer) { ... }
    TfLiteStatus Prepare(TfLiteOpaqueContext* context,
                         TfLiteOpaqueNode* node) { ... }
    TfLiteStatus Invoke(TfLiteOpaqueContext* context,
                        TfLiteOpaqueNode* node) {... }
  };

  const TfLiteOperator* MyCustomOperator() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteOperator* my_custom_op = ()[] {
        TfLiteOperator* r =
            TfLiteOperatorCreate(
                kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
        TfLiteOperatorSetInit(r, Init);
        TfLiteOperatorSetFree(r, Free);
        TfLiteOperatorSetPrepare(r, Prepare);
        TfLiteOperatorSetInvoke(r, Eval);
        return r;
      };
    return my_custom_op;
  }
}  // namespace my_namespace

C

static void* MyCustomOpInit(TfLiteOpaqueContext* context, const char* buffer,
                     size_t length) { ... }
static void MyCustomOpFree(TfLiteOpaqueContext* context, void* buffer) { ... }
static TfLiteStatus MyCustomOpPrepare(TfLiteOpaqueContext* context,
                                      TfLiteOpaqueNode* node) { ... }
static TfLiteStatus MyCustomOpInvoke(TfLiteOpaqueContext* context,
                                     TfLiteOpaqueNode* node) {... }

static TfLiteOperator* MyCustomOpCreate() {
  const TfLiteOperator* r =
      TfLiteOperatorCreate(
          kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
  TfLiteOperatorSetInit(r, MyCustomOpInit);
  TfLiteOperatorSetFree(r, MyCustomOpFree);
  TfLiteOperatorSetPrepare(r, MyCustomOpPrepare);
  TfLiteOperatorSetInvoke(r, MyCustomOpEval);
  return r;
}

const TfLiteOperator* MyCustomOperator() {
  // Singleton instance, intentionally never destroyed.
  static const TfLiteOperator* my_custom_op = MyCustomOpCreate();
  return my_custom_op;
}

登録は自動ではなく、MyCustomOperator 関数への明示的な呼び出しを行う必要があります（詳細は下記を参照）。標準の BuiltinOpResolver（:builtin_ops ターゲットから利用可能）は組み込みの登録を処理しますが、カスタムオペレーションは個別のカスタムライブラリで収集する必要があります。

LiteRT ランタイムでのカーネルの定義

LiteRT で op を使用するには、2 つの関数（Prepare と Eval）と、TfLiteOperator を構築する 3 つ目の関数を定義するだけです。

C++

namespace atan_op {
  namespace {
    TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
      TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1);
      TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1);

      const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
      TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

      int num_dims = TfLiteOpaqueTensorNumDimensions(input);

      TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
      for (int i=0; i < num_dims; ++i) {
        output_size->data[i] = input->dims->data[i];
      }

      return TfLiteOpaqueContextResizeTensor(context, output, output_size);
    }

    TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
      const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
      TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

      float* input_data = static_cast<float*>(TfLiteOpaqueTensorData(input));
      float* output_data = static_cast<float*>(TfLiteOpaqueTensorData(output));

      size_t count = 1;
      int num_dims = TfLiteOpaqueTensorNumDimensions(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;
    }
  }  // anonymous namespace

  const TfLiteOperator* AtanOperator() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteOperator* atan_op = ()[] {
        auto* r = TfLiteOperatorCreate(
            kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
        TfLiteOperatorSetPrepare(r, Prepare);
        TfLiteOperatorSetInvoke(r, Eval);
        return r;
      };
    return atan_op;
  }
}  // namespace atan_op

C

static TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
  TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1);
  TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1);

  const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
  TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

  int num_dims = TfLiteOpaqueTensorNumDimensions(input);

  TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
  for (int i = 0; i < num_dims; ++i) {
    output_size->data[i] = input->dims->data[i];
  }

  return TfLiteOpaqueContextResizeTensor(context, output, output_size);
}

static TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
  const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
  TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

  float* input_data = static_cast<float*>(TfLiteOpaqueTensorData(input));
  float* output_data = static_cast<float*>(TfLiteOpaqueTensorData(output));

  size_t count = 1;
  int num_dims = TfLiteOpaqueTensorNumDimensions(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;
}

static const TfLiteOperator* AtanOpCreate() {
  TfLiteOperator* r = TfLiteOperatorCreate(
          kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
  TfLiteOperatorSetPrepare(r, Prepare);
  TfLiteOperatorSetInvoke(r, Eval);
  return r;
}

const TfLiteOperator* AtanOperator() {
  // Singleton instance, intentionally never destroyed.
  static const TfLiteOperator* atan_op = AtanOpCreate();
  return atan_op;
}

OpResolver を初期化するときに、カスタムオペレーションをリゾルバに追加します（例については下記を参照）。これにより、LiteRT にオペレーターが登録され、LiteRT で新しい実装を使用できるようになります。

演算子をカーネルライブラリに登録する

次に、演算子をカーネルライブラリに登録する必要があります。これには OpResolver を使用します。バックグラウンドでは、インタープリタがカーネルのライブラリを読み込み、モデル内の各演算子を実行するように割り当てます。デフォルトのライブラリには組み込みカーネルのみが含まれていますが、カスタムライブラリ op 演算子で置き換えたり、拡張したりすることは可能です。

演算子コードと名前を実際のコードに変換する OpResolver クラスは、次のように定義されます。

class OpResolver {
 public:
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  ...
};

下位互換性のため、このクラスでは不透明型 TfLiteOperator ではなく古い具象型 TfLiteRegistration が使用されていますが、TfLiteRegistration 構造体には TfLiteOperator* 型の registration_external フィールドが含まれています。

MutableOpResolver クラスと BuiltinOpResolver クラスは OpResolver から派生します。

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddAll(const MutableOpResolver& other);
  ...
};

class BuiltinOpResolver : public MutableOpResolver {
 public:
  BuiltinOpResolver();  // Constructs an op resolver with all the builtin ops.
};

通常の使用（カスタムオペレーションなし）では、BuiltinOpResolver を使用して次のように記述する必要があります。

tflite::ops::builtin::BuiltinOpResolver resolver;

上記で作成したカスタムオペレーションを追加するには、代わりに MutableOpResolver を使用し、（リゾルバを InterpreterBuilder に渡す前に）tflite::AddOp を呼び出すことができます。

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
tflite::AddOp(&resolver, AtanOpRegistration());

組み込みオペレーションのセットが大きすぎると判断された場合は、指定されたオペレーションのサブセット（指定されたモデルに含まれるオペレーションのみなど）に基づいて、新しい OpResolver をコード生成できます。これは TensorFlow の選択的登録と同等です（その簡単なバージョンは tools ディレクトリにあります）。

Java でカスタムオペレータを定義する場合は、現在、独自のカスタム JNI レイヤを構築し、この JNI コードで独自の AAR をコンパイルする必要があります。同様に、Python で使用可能なこれらの演算子を定義する場合は、Python ラッパーコードに登録を配置できます。

上記のプロセスと同様のプロセスは、単一のオペレーターではなく一連のオペレーションをサポートする場合にも使用できます。必要な数の AddCustom 演算子を追加するだけです。また、MutableOpResolver を使用すると、AddBuiltin を使用して組み込みの実装をオーバーライドすることもできます。

演算子をテストしてプロファイリングする

LiteRT ベンチマークツールで op のプロファイルを作成するには、LiteRT のベンチマークモデルツールを使用します。テスト目的で、適切な AddCustom 呼び出し（上記参照）を register.cc に追加することで、LiteRT のローカルビルドでカスタムオペレーションを認識させることができます。

ベストプラクティス

メモリの割り当てと割り当て解除は慎重に最適化してください。Prepare でメモリを割り当てる方が Invoke で割り当てるよりも効率的です。また、ループの前にメモリを割り当てる方が、各イテレーションで割り当てるよりも効率的です。自分で malloc するのではなく、一時テンソルデータを使用します（項目 2 を参照）。可能な限りコピーではなくポインタ/リファレンスを使用します。

データ構造がオペレーション全体で存続する場合は、一時テンソルを使用してメモリを事前割り当てすることをおすすめします。他の関数でテンソルインデックスを参照するには、OpData 構造体を使用する必要がある場合があります。畳み込みのカーネルの例をご覧ください。サンプルコードスニペットを以下に示します。

struct MyOpData {
  int temp_tensor_index;
  ...
};

void* Init(TfLiteOpaqueContext* context,
    const char* buffer, size_t length) {
  auto* op_data = new MyOpData{};
  ...
  return op_data;
}
void Free(TfLiteOpaqueContext* context, void* buffer) {
  ...
  delete reinterpret_cast<MyOpData*>(buffer);
}
TfLiteStatus Prepare(TfLiteOpaqueContext* context,
                     TfLiteOpaqueNode* node) {
  ...
  auto* op_data =
      reinterpret_cast<MyOpData*>(TfLiteOpaqueNodeGetUserData(node));
  const int num_temporaries = 1;
  int temporary_tensor_indices[num_temporaries];
  TfLiteOpaqueTensorBuilder* builder = TfLiteOpaqueTensorBuilderCreate();
  TfLiteOpaqueTensorBuilderSetType(builder, kTfLiteFloat32);
  TfLiteOpaqueTensorBuilderSetAllocationType(builder, kTfLiteArenaRw);
  TfLiteOpaqueContextAddTensor(context, builder,
      &temporary_tensor_indices[0]);
  TfLiteOpaqueTensorBuilderDelete(builder);
  TfLiteOpaqueNodeSetTemporaries(node, temporary_tensor_indices,
      num_temporaries);
  op_data->temp_tensor_index = temporary_tensor_indices[0];
  ...
  return kTfLiteOk;
}
TfLiteStatus Invoke(TfLiteOpaqueContext* context,
                    TfLiteOpaqueNode* node) {
  ...
  auto* op_data = reinterpret_cast<MyOpData*>(
      TfLiteOpaqueNodeGetUserData(node));
  TfLiteOpaqueTensor* temp_tensor =
      TfLiteOpaqueContextGetOpaqueTensor(context,
          op_data->temp_tensor_index);
  TF_LITE_OPAQUE_ENSURE(context,
      TfLiteTensorType(temp_tensor) == kTfLiteFloat32);
  TF_LITE_OPAQUE_ENSURE(context,
      TfLiteTensorGetAllocationType(temp_Tensor) == kTfLiteArenaRw);
  void *temp_data = TfLiteTensorData(temp_tensor);
  TF_LITE_OPAQUE_ENSURE(context, temp_data != nullptr);
  ...
  return kTfLiteOk;
}

メモリの無駄遣いがそれほど多くない場合は、実行のたびに動的に割り当てられる std::vector よりも、静的な固定サイズの配列（または Resize の事前割り当てられた std::vector）を使用することをおすすめします。
バイナリサイズに影響するため、まだ存在しない標準ライブラリコンテナテンプレートのインスタンス化は避けてください。たとえば、他のカーネルに存在しない std::map がオペレーションに必要な場合、直接インデックスマッピングで std::vector を使用すると、バイナリサイズを小さく保ちながら、そのオペレーションを実行できます。他のカーネルが使用しているものを確認して、分析情報を取得します（または質問します）。
malloc が返すメモリへのポインタを確認します。このポインタが nullptr の場合、そのポインタを使用して操作を行うべきではありません。関数で malloc を使用し、エラー終了がある場合は、終了前にメモリを割り当て解除します。
TF_LITE_OPAQUE_ENSURE(context, condition) を使用して、特定の条件を確認します。コードでは、TF_LITE_OPAQUE_ENSURE を使用したときにメモリがハングアップしないようにする必要があります。つまり、これらのマクロは、リークするリソースが割り当てられる前に使用する必要があります。