Google Edge TPUで、TensorFlow Liteを使う例！

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

Coralが公開しているedgetpu-nativeから、Google Edge TPUで、TensorFlow Liteを使う例のソースコードを覗いてみます。

edgetpu/cpp/examples/utils.ccのBuildEdgeTpuInterpreter

std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
    const tflite::FlatBufferModel& model,
    edgetpu::EdgeTpuContext* edgetpu_context) {
  tflite::ops::builtin::BuiltinOpResolver resolver;
  resolver.AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp());
  std::unique_ptr<tflite::Interpreter> interpreter;
  if (tflite::InterpreterBuilder(model, resolver)(&interpreter) != kTfLiteOk) {
    std::cerr << "Failed to build interpreter." << std::endl;
  }
  // Bind given context with interpreter.
  interpreter->SetExternalContext(kTfLiteEdgeTpuContext, edgetpu_context);
  interpreter->SetNumThreads(1);
  if (interpreter->AllocateTensors() != kTfLiteOk) {
    std::cerr << "Failed to allocate tensors." << std::endl;
  }
  return interpreter;
}

resolver.AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp())

にて、カスタムOpとして、libedgetpu/edgetpu.hにあるedgetpu-custom-opを登録しているようです。

この edgetpu-custom-op というOpって、6月3日のGoogle Edge TPUの offline compilerで出てきた Offline Compilerのログに subgraph とあったものをひとまとまりにしたものかな？

TensorFlow XLAでやっているのがそのような方式なので。。。

そうじゃなければ、どんな仕組みなんだろうか？

裏付け、見つけました。@PINTO03091さんのこのツイート

ここにあるのは、Subgraph0の中身はCUSTOMなOpとそのOpに対する入力と出力だけです。

ということで、TensorFlow XLAと同じ感じに、1つのOpにまとめて、中でなんかやっているようです。

なお、kTfLiteEdgeTpuContext は、TensorFlow Liteのソースコードでちゃんと定義されています。