在 Tensorflow Lite C API 中注冊自定義運算符

Question

我正在使用 C API 在 Android 上運行 tensorflow lite。 我的模型需要操作符RandomStandardNormal ，它最近在 tensorflow v2.4.0-rc0 here 中作為自定義操作原型實現

TfLiteInterpreterOptionsAddCustomOp()函數在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);

看看這個例子和線程，我試圖像這樣使用TfLiteInterpreterOptionsAddCustomOp ：

// 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);

我看到Register_RANDOM_STANDARD_NORMAL()函數是在tensorflow/lite/kernels/custom_ops_register.h的tflite::ops::custom C++ 命名空間中tflite::ops::custom 。 但是，當我嘗試將它包含在我的 C 文件中時，編譯器會抱怨，因為namespace是 C 中的未知類型。

如何使用 tensorflow lite C API 注冊自定義運算符？ 我是否需要使用 C++ 編譯器才能將 C API 與此自定義運算符一起使用，因為它是在 C++ 中定義的？

注意：我在編譯libtensorflowlite_c.so時在 bazel BUILD deps 中包含//tensorflow/lite/kernels:custom_ops libtensorflowlite_c.so

Answer 1

看起來這是通過以下解決方法在 Github 上回答的：

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

Answer 2

在 tensorflow github 上，@jdduke 提出了一個臨時解決方法：

• 向custom_ops_register.h添加一個extern "C"包裝頭

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

• 向random_standard_normal.cc添加extern "C"包裝器實現

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

• 確保//tensorflow/lite/kernels:custom_ops作為依賴項包含在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
    ],
)

• 修改我的 C++ 代碼以調用這個新的包裝函數

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

它奏效了！ 我的張量終於在 android 上分配了 :)