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
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