c++中載入tensorflow serving模型格式檔案
阿新 • • 發佈:2019-01-04
前幾篇文章在講c++中載入pb格式檔案,就是單純的pb,沒有變數的情況,下午仔細看了下c++的原始碼發現是可以直接載入tensorflow serving格式檔案,格式檔案包括一個pb檔案和一些variables變數資料夾,廢話不多說,直接看程式碼:
CMakeLists.txt:
cmake_minimum_required(VERSION 3.10) project(cppexcise) set(CMAKE_CXX_STANDARD 11) link_directories(/Users/xxxx/Documents/tensorflow/bazel-bin/tensorflow) include_directories( /Users/xxxx/Documents/tensorflow /Users/xxxx/Documents/tensorflow/bazel-genfiles /Users/xxxx/Documents/tensorflow/bazel-bin/tensorflow /Users/xxxx/Downloads/eigen3) add_executable(cppexcise main.cpp ) target_link_libraries(cppexcise tensorflow_cc tensorflow_framework)
c++簡單程式碼:
#include <iostream> #include <vector> #include "tensorflow/cc/saved_model/loader.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/protobuf/meta_graph.pb.h" #include "tensorflow/cc/saved_model/tag_constants.h" using namespace std; using namespace tensorflow; int main(int argc ,char *argv[]) { string modelpath; if(argc<2){ cout<<"請輸入模型路徑"; return 0; }else{ modelpath=argv[1]; } tensorflow::SessionOptions sess_options; tensorflow::RunOptions run_options; tensorflow::SavedModelBundle bundle; Status status; status =tensorflow::LoadSavedModel(sess_options, run_options, modelpath, {tensorflow::kSavedModelTagServe}, &bundle); if(!status.ok()){ cout<<status.ToString()<<endl; } tensorflow::MetaGraphDef graph_def = bundle.meta_graph_def; std::unique_ptr<tensorflow::Session>& session = bundle.session; vector<int> vec={7997, 1945, 8471, 14127, 17565, 7340, 20224, 17529, 3796, 16033, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; int ndim=vec.size(); Tensor x(tensorflow::DT_INT32, tensorflow::TensorShape({1, ndim})); // New Tensor shape [1, ndim] auto x_map = x.tensor<int, 2>(); for (int j = 0; j < ndim; j++) { x_map(0, j) = vec[j]; } std::vector<std::pair<string, tensorflow::Tensor>> inputs; inputs.push_back(std::pair<std::string, tensorflow::Tensor>("input_x", x)); Tensor keep_prob(tensorflow::DT_FLOAT, tensorflow::TensorShape({1})); keep_prob.vec<float>()(0) = 1.0f; inputs.push_back(std::pair<std::string, tensorflow::Tensor>("keep_prob", keep_prob)); Tensor tensor_out(tensorflow::DT_INT32, TensorShape({1,ndim})); std::vector<tensorflow::Tensor> outputs={{ tensor_out }}; status= session->Run(inputs, {"crf_pred/ReverseSequence_1"}, {}, &outputs); if (!status.ok()) { std::cout << status.ToString() << "\n"; return 1; } for(int i=0;i<40;++i) { std::cout << outputs[0].matrix<int>()(0,i)<<" "; } cout<<endl; return 0; }