在 Tensorflow Lite C API 中注册自定义运算符

Question

I am running tensorflow lite on Android using the C API. My model requires the operator RandomStandardNormal which was recently implemented as a custom op prototype in tensorflow v2.4.0-rc0 here

TfLiteInterpreterOptionsAddCustomOp() function is listed in tensorflow/lite/c/c_api_experimental.h :

TFL_CAPI_EXPORT void TfLiteInterpreterOptionsAddCustomOp(
    TfLiteInterpreterOptions* options, const char* name,
    const TfLiteRegistration* registration, int32_t min_version,
    int32_t max_version);

Looking at this example & thread , I am trying to use TfLiteInterpreterOptionsAddCustomOp like this:

// create model and interpreter options
TfLiteModel *model = TfLiteModelCreateFromFile("path/to/model.tflite");
TfLiteInterpreterOptions* options = TfLiteInterpreterOptionsCreate();

// register custom ops
TfLiteInterpreterOptionsAddCustomOp(options, "RandomStandardNormal", Register_RANDOM_STANDARD_NORMAL(), 1, 1);

// create the interpreter
TfLiteInterpreter* interpreter = TfLiteInterpreterCreate(model, options);
TfLiteInterpreterAllocateTensors(interpreter);

I see that the Register_RANDOM_STANDARD_NORMAL() function is defined in the tflite::ops::custom C++ namespace in tensorflow/lite/kernels/custom_ops_register.h . But, when I try to include this in my C file the compiler complains because namespace is an unknown type in C.

How can I register a custom operator using the tensorflow lite C API? Do I need to use a C++ compiler in order to use the C API with this custom operator because it was defined in C++?

NOTE: I include //tensorflow/lite/kernels:custom_ops in the bazel BUILD deps when compiling libtensorflowlite_c.so

Answer 1

It looks like this was answered on Github via this workaround:

https://github.com/tensorflow/tensorflow/issues/44664#issuecomment-723310060

Answer 2

On tensorflow github , @jdduke suggested a temporary workaround:

• add a extern "C" wrapper head to custom_ops_register.h

extern "C" {
TFL_CAPI_EXPORT TfLiteRegistration* TfLiteRegisterRandomStandardNormal();
}

• add a extern "C" wrapper implementation to random_standard_normal.cc

extern "C" {
TFL_CAPI_EXPORT TfLiteRegistration* TfLiteRegisterRandomStandardNormal() {
  return tflite::ops::custom::Register_RANDOM_STANDARD_NORMAL();
}
}

• ensure //tensorflow/lite/kernels:custom_ops is included as a dependency in tensorflow/lite/c/BUILD

tflite_cc_shared_object(
    name = "tensorflowlite_c",
    linkopts = select({
        "//tensorflow:ios": [
            "-Wl,-exported_symbols_list,$(location //tensorflow/lite/c:exported_symbols.lds)",
        ],
        "//tensorflow:macos": [
            "-Wl,-exported_symbols_list,$(location //tensorflow/lite/c:exported_symbols.lds)",
        ],
        "//tensorflow:windows": [],
        "//conditions:default": [
            "-z defs",
            "-Wl,--version-script,$(location //tensorflow/lite/c:version_script.lds)",
        ],
    }),
    per_os_targets = True,
    deps = [
        ":c_api",
        ":c_api_experimental",
        ":exported_symbols.lds",
        ":version_script.lds",
        "//tensorflow/lite/kernels:custom_ops", # here
    ],
)

• modify my C++ code to call this new wrapper function

TfLiteInterpreterOptionsAddCustomOp(options, "RandomStandardNormal", TfLiteRegisterRandomStandardNormal(), 1, 1);

And it worked! My tensors finally allocated on android :)