c++中加载tensorflow serving模型格式文件

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/luoyexuge/article/details/81874325

前几篇文章在讲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;

}