1. 程式人生 > >TensorFlow模型儲存的一個坑

TensorFlow模型儲存的一個坑

最近工作需要,使用tensorflow(r1.3)中contrib的某個庫,根據tf官網的說法,go版本是隻能呼叫py版本訓練好的模型,而個人對go比較熟且線上環境用go作服務,所以想用go調python訓練好的模型。

使用python import tf和tf.contrib相關庫,使用tf.saved_model.builder.SavedModelBuilder儲存模型,但使用go使用LoadSavedModel的時候,會報形如Not found Op type not registered 'CountExtremelyRandomStats' 的錯誤,具體的話是說在libtensorflow裡面找不到CountExtremelyRandomStats

這個op。

嘗試了很多方案,略過。最終發現問題是TensorFlow在BUILD target:libtensorflow.so的時候,根本沒有把contrib相關內容打進去。修改根BUILD檔案的libtensorflow.so部分如下:

cc_binary(
    name = "libtensorflow.so",
    linkopts = select({
        "//tensorflow:darwin": [
            "-Wl,-exported_symbols_list",  # This line must be directly followed by the exported_symbols.lds file
"//tensorflow/c:exported_symbols.lds", ], "//tensorflow:windows": [], "//tensorflow:windows_msvc": [], "//conditions:default": [ "-z defs", "-s", "-Wl,--version-script", # This line must be directly followed by the version_script.lds file
"//tensorflow/c:version_script.lds", ], }), linkshared = 1, deps = [ "//tensorflow/contrib:contrib_kernels", "//tensorflow/contrib:contrib_ops_op_lib", "//tensorflow/c:c_api", "//tensorflow/c:exported_symbols.lds", "//tensorflow/c:version_script.lds", "//tensorflow/core:tensorflow", ], )

contrib那兩行依賴是新加進去的,然後重新用bazel編譯一個libtensorflow就可以了。期間困惑的是為什麼python用pip裝的就可以用contrib呢,觀察tools/pip_packages/BUILD檔案發現pip的包,是包含了所有contrib的依賴的。感覺官方團隊對非python語言的重視還不夠,沒有明顯說明這個東西,有人提了issue也沒有人迴應(我最後迴應了這個op相關的一個issue:https://github.com/tensorflow/tensorflow/issues/14358)。