TensorFlow Liteのカスタムオペレーション

1. はじめに

「TensorFlow Lite」は、TensorFlowのオペレーションをすべてサポートしているわけではありません。モデルにサポートされていないオペレーションが含まれている時は、ユーザー実装のC/C++コードを「カスタムオペレーション」として使用できます。
(情報源)

2. カスタムオペレーションの使用手順

カスタムオペレーションの使用手順は次の通りです。

(1)グラフ内で正しい名前のオペレーションを参照していることを確認。
(2)「カスタムオペレーション」を「TensorFlow Lite」に登録。
ランタイムがグラフ内のオペレーションをカスタムオペレーションにマップできるようになります。
(3)オペレーションの正確さとパフォーマンスをプロファイリング。
カスタムオペレーションのみをテストする場合は、カスタムオペレーションのみで、「benchmark_model」を使用してモデルを作成することをお勧めします。

3. Sinのカスタムオペレーションの作成

「TensorFlow Lite」にはないオペレーションをサポートする例を見てみましょう。Sinオペレーションを使用して、関数「y = sin(x + offset)」の非常に単純なモデルを構築していると仮定します。ここで、offsetは訓練可能です。

◎TensorFlowモデルの生成
TensorFlowモデルを訓練するコードは、次のようになります。

offset = tf.get_variable("offset", [1,], tf.float32)
x = tf.placeholder(tf.float32, shape=(None,))
y = tf.sin(x + offset)
y_ = tf.placeholder(tf.float32, shape=(None,))
loss = tf.reduce_sum(tf.square(y - y_))
optimizer = tf.train.GradientDescentOptimizer(0.001)
train = optimizer.minimize(loss)

◎TensorFlow Liteモデルへの変換と実行
「--allow_custom_ops」引数を指定して「TensorFlow Lite Optimizing Converter」を使用して、このモデルをTensorflow Lite形式に変換し、デフォルトのインタープリタで実行すると、次のエラーメッセージが表示されます。

Didn't find custom op for name 'Sin'
Registration failed.

◎TensorFlow Liteランタイムのカスタムオペレーションの定義
「TensorFlow Lite」で「カスタムオペレーション」を使用するために必要なことは、「Prepare」「Eval」の2つの関数を定義し、「TfLiteRegistration」を構築することのみです。コードは次のようになります。

TfLiteStatus SinPrepare(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 SinEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node,0);
  TfLiteTensor* output = GetOutput(context, node,0);

  float* input_data = input->data.f;
  float* output_data = output->data.f;

  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] = sin(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_SIN() {
  static TfLiteRegistration r = {nullptr, nullptr, SinPrepare, SinEval};
  return &r;
}

OpResolverを初期化するときに、「カスタムオペレーション」を追加します。これにより、「TensorFlow Lite」に「カスタムオペレーション」が追加され、使用できるようになります。

「TfLiteRegistration」の最後の2つの引数は、「カスタムオペレーション」に対して定義したSinPrepare()とSinEval()に対応します。オペレーションで使用される変数を初期化し、スペースを解放するために2つの関数を使用した場合、Init()とFree()は、TfLiteRegistrationの最初の2つの引数に追加されます。この例ではnullptrに設定されています。

tflite::ops::builtin::BuiltinOpResolver builtins;
builtins.AddCustom("Sin", Register_SIN());

Javaでカスタムオペレーションを作成する場合、現在、独自のカスタムJNIレイヤーを構築し、このJNIコードで独自のAARをコンパイルする必要があります。

同様に、これらのオペレーションをPythonで使用できるようにする場合は、Pythonラッパーコードに登録を配置できます。

単一のオペレーションではなく、一連の操作をサポートするために、上記と同様の手順で対応できます。必要な数のAddCustomオペレーションを追加するだけです。さらに、「BuiltinOpResolver」では、「AddBuiltin」を使用して組み込みの実装をオーバーライドすることもできます。

iPhone/Androidアプリ開発者のための機械学習・深層学習 実践入門