tflite添加自定义op

1 添加自定义 op 的文件

文件：tensorflow/lite/kernels/crop_and_resize.cc

namespace tflite {
namespace ops {
namespace custom {
namespace crop_and_resize {

constexpr int kInputTensor = 0;
constexpr int kBoxesTensor = 1;
constexpr int kBoxIndTensor = 2;
constexpr int kSizeTensor = 3;
constexpr int kOutputTensor = 0;

...

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node)
{
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 4);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

...

}

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node)
{

...

  return kTfLiteOk;
}
}  // namespace crop_and_resize

TfLiteRegistration* Register_CROP_AND_RESIZE() {
  static TfLiteRegistration r = {nullptr, nullptr,
                                crop_and_resize::Prepare,
                                crop_and_resize::Eval};
  return &r;
}

}  // namespace builtin
}  // namespace ops
}  // namespace tflite

2 注册自定义 op

文件：tensorflow/lite/kernels/register.cc

namespace custom {

TfLiteRegistration* Register_AUDIO_SPECTROGRAM();
TfLiteRegistration* Register_LAYER_NORM_LSTM();
TfLiteRegistration* Register_MFCC();
TfLiteRegistration* Register_DETECTION_POSTPROCESS();
TfLiteRegistration* Register_RELU_1();
TfLiteRegistration* Register_CROP_AND_RESIZE();

}  // namespace custom

...
...
...

  // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that
  // custom ops aren't always included by default.
  AddCustom("Mfcc", tflite::ops::custom::Register_MFCC());
  AddCustom("AudioSpectrogram",
            tflite::ops::custom::Register_AUDIO_SPECTROGRAM());
  AddCustom("LayerNormLstm", tflite::ops::custom::Register_LAYER_NORM_LSTM());
  AddCustom("Relu1", tflite::ops::custom::Register_RELU_1());
  AddCustom("TFLite_Detection_PostProcess",
            tflite::ops::custom::Register_DETECTION_POSTPROCESS());
  AddCustom("Crop_and_Resize",
            tflite::ops::custom::Register_CROP_AND_RESIZE());
}

3 修改编译文件

文件：tensorflow/lite/kernels/BUILD

cc_library(
    name = "builtin_op_kernels",
    srcs = [
        "activations.cc",
        "add.cc",
        "arg_min_max.cc",

...

        "crop_and_resize.cc",
],

4 按照官方教程 源码编译

4.1 安装bazel

bazel官网安装教程

4.2 配置编译系统

cd tensorflow
./configure

4.3 Bazel build

支持GPU

bazel build --config=opt --config=cuda //tensorflow/tools/pip_package:build_pip_package

4.4 编译软件包

./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

4.5 安装软件包

pip3 install /tmp/tensorflow_pkg/tensorflow-version-tags.whl --user