C++加载PB模型文件

C++加载PB模型

• OS：Ubantu18.04
• Tensorflow动态库版本：Tensorflow2.1
• OpenCV版本：3.7
• cmake：3.10.2

MakeLists.txt

cmake_minimum_required(VERSION 2.8.8)
project(demo)

set(CMAKE_CXX_STANDARD 11)
set(tensorflow_path /home/tensorflow)
set(current_project /home/yuxiubin-test/demo/)

# Find OpenCV, you may need to set OpenCV_DIR variable
# to the absolute path to the directory containing OpenCVConfig.cmake file
# via the command line or GUI
find_package(OpenCV REQUIRED)

# If the package has been found, several variables will
# be set, you can find the full list with descriptions
# in the OpenCVConfig.cmake file.
# Print some message showing some of them
message(STATUS "OpenCV library status:")
message(STATUS "version: ${OpenCV_VERSION}")
message(STATUS "libraries: ${OpenCV_LIBS}")
message(STATUS "include path: ${OpenCV_INCLUDE_DIRS}")

include_directories(${current_project})
include_directories(${tensorflow_path})
include_directories(${tensorflow_path}/bazel-genfiles)
include_directories(${tensorflow_path}/bazel-bin/tensorflow)
include_directories(${tensorflow_path}/tensorfow/contrib/makefile/gen/proto)
include_directories(${tensorflow_path}/tensorflow/contrib/makefile/gen/protobuf-host/include)
include_directories(${tensorflow_path}/tensorflow/contrib/makefile/downloads/eigen)
include_directories(${tensorflow_path}/tensorflow/contrib/makefile/downloads/absl)
include_directories(${tensorflow_path}/tensorflow/contrib/makefile/downloads/nsync/public)
#设置链接库搜索目录
link_directories(${tensorflow_path}/bazel-bin/tensorflow)

# Add OpenCV headers location to your include paths
#设置包含的目录
include_directories(${OpenCV_INCLUDE_DIRS})

#新建demo 其源码为Main.cpp
add_executable(demo Main.cpp) 

# 添加链接库
target_link_libraries(demo tensorflow_cc tensorflow_framework ${OpenCV_LIBS})

主要有三个步骤：

1. 加载保存的pb模型文件，生成graph和session
2. 使用OpenCV读取图像进行处理，将OpenCV的Mat转成适合模型读入的tensor类型
3. 放入网络sess.Run()方法获取输出，前提要知道输入输出层的名称。

第一步：加载pb模型文件，主要使用了LoadSavedModel这一个方法，该方法会返回SavedModelBundle对象，通过该对象可以获取graph和session，具体步骤如下：

	tensorflow::SessionOptions sess_options;
    tensorflow::RunOptions run_options;
    tensorflow::SavedModelBundle bundle;
	Status status_create=LoadSavedModel(sess_options,run_options,model_path, {tensorflow::kSavedModelTagServe}, &bundle);
	if(!status_create.ok()){
		cout << "ERROR: Creating graph in session failed..." << status_create.ToString() << std::endl;
        return -1;
	}
    cout << "<----Successfully created session and load graph.------->"<< endl;
	
	/*--------------------------------创建session------------------------------*/
    tensorflow::MetaGraphDef graph_def = bundle.meta_graph_def;
    std::unique_ptr<tensorflow::Session>& session = bundle.session;//创建新会话Session

第二步：使用OpenCV读取图片并进行灰度化处理，并将Mat转成四维的tensor，作为模型的输出，这里mat转成tensor的方法是一个技巧，利用了指针。

    Mat imagePixels=imread(image_path);
    if(imagePixels.empty())
    {
        cout<<"can't open the image!!!!!!!"<<endl;
        return -1;
    }
	//转化为灰度图
	cvtColor(imagePixels,imagePixels,CV_BGR2GRAY);
	
    //创建一个tensor作为输入网络的接叿    
    //将Opencv的Mat格式的图片存入tensor
	//四维的输入tensor
    Tensor input_tensor(DT_FLOAT, TensorShape({1,input_height,input_width,1}));

	// 指针指向输入tensor的内存
    float *p = input_tensor.flat<float>().data();
	// create a "fake" cv::Mat from it 
    Mat img(input_height, input_width, CV_32FC1, p);
    // use it here as a destination
    cout << "Image Pixels: " << imagePixels.empty() << " " << imagePixels.cols << "x" << imagePixels.rows << endl;

    //做归一化opecv c语言中的normalize和python中的参数不太一样需要注意    
    //normalize(imagePixels, imagePixels, 1.0, 0.0, NORM_L2, CV_32F);
    //存入img的同时，也将数据存放到了输入tensor
	imagePixels.convertTo(img, CV_32FC1);

第三步：执行sess.Run()方法，执行该方法之前需要知道模型的输入输出层的名称是什么，才能向模型传入输入，获取模型的输出，可以使用saved_model_cli工具查看。

    string input_tensor_name="serving_default_inputs:0";        //模型的输入层名称
	string output_tensor_name="StatefulPartitionedCall:";       //模型的输出层名称
	//前向运行，输出结果一定是一个tensor的vector
	vector<tensorflow::Tensor> outputs;
	char result[5];
	//因为有五个输出所以循环了五次
    for(int i=0;i<5;i++)
	{
		 
		string output_node = output_tensor_name+to_string(i);
		Status status_run = session->Run({{input_tensor_name, input_tensor}}, {output_node}, {}, &outputs);
 
		if (!status_run.ok()) {
			cout << "ERROR: RUN failed..."  << std::endl;
			cout << status_run.ToString() << "\n";
			return -1;
		}
		int class_id=GetLabelFromOutput(outputs);
		result[i]=INDEX2CH[class_id];
		outputs.clear();
	}

Main.cpp

#include <fstream>
#include <utility>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <iostream>
#include<opencv2/opencv.hpp>
 
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
 
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
 
#include "tensorflow/core/graph/default_device.h"
#include "tensorflow/core/graph/graph_def_builder.h"
 
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
 
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/command_line_flags.h"
 
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
 
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include  "tensorflow/cc/saved_model/tag_constants.h"
#include "opencv2/opencv.hpp"

#include "class_dict.h"
using namespace tensorflow::ops;
using namespace tensorflow;
using namespace std;
using namespace cv;
using tensorflow::Flag;
using tensorflow::Tensor;
using tensorflow::Status;
using tensorflow::string;
using tensorflow::int32 ;
int GetLabelFromOutput(vector<tensorflow::Tensor> outputs)
{
	//把输出值给提取出来
	
    Tensor t = outputs[0];                   // Fetch the first tensor
    auto tmap = t.tensor<float, 2>();        // Tensor Shape: [batch_size, target_class_num]
    int output_dim = t.shape().dim_size(1);  // Get the target_class_num from 1st dimension
 
    // Argmax: Get Final Prediction Label and Probability
    int output_class_id = -1;
    double output_prob = 0.0;
    for (int j = 0; j < output_dim; j++)
    {
        if (tmap(0, j) >= output_prob) {
            output_class_id = j;
            output_prob = tmap(0, j);
        }
    }
 
	return output_class_id;
}

int main(int argc, char** argv )
{
    /*--------------------------------配置关键信息------------------------------*/
    string model_path="/home/yuxiubin-test/demo/model_tf2V2/";
    string image_path;
	cout<<"请输入图片路径"<<endl;
	cin>>image_path;
    int input_height =40;
    int input_width=100;
    string input_tensor_name="serving_default_inputs:0";
	string output_tensor_name="StatefulPartitionedCall:";
	
	
    tensorflow::SessionOptions sess_options;
    tensorflow::RunOptions run_options;
    tensorflow::SavedModelBundle bundle;
 
    
    /*--------------------------------从pb文件中读取模型-------------------------------*/
 
	Status status_create=LoadSavedModel(sess_options,run_options,model_path, {tensorflow::kSavedModelTagServe}, &bundle);
	if(!status_create.ok()){
		cout << "ERROR: Creating graph in session failed..." << status_create.ToString() << std::endl;
        return -1;
	}
    cout << "<----Successfully created session and load graph.------->"<< endl;
	
	/*--------------------------------创建session------------------------------*/
    tensorflow::MetaGraphDef graph_def = bundle.meta_graph_def;
    std::unique_ptr<tensorflow::Session>& session = bundle.session;//创建新会话Session
 
    /*---------------------------------载入测试图片-------------------------------------*/
    cout<<endl<<"<------------loading test_image-------------->"<<endl;
    Mat imagePixels=imread(image_path);
    if(imagePixels.empty())
    {
        cout<<"can't open the image!!!!!!!"<<endl;
        return -1;
    }
	//转化为灰度图
	cvtColor(imagePixels,imagePixels,CV_BGR2GRAY);
	
    //创建一个tensor作为输入网络的接叿    
    //将Opencv的Mat格式的图片存入tensor
	//四维的输入tensor
    Tensor input_tensor(DT_FLOAT, TensorShape({1,input_height,input_width,1}));

	// get pointer to memory for that Tensor
    float *p = input_tensor.flat<float>().data();
	// create a "fake" cv::Mat from it 
    Mat img(input_height, input_width, CV_32FC1, p);
    // use it here as a destination
    cout << "Image Pixels: " << imagePixels.empty() << " " << imagePixels.cols << "x" << imagePixels.rows << endl;

    //做归一化opecv c语言中的normalize和python中的参数不太一样需要注意    
    //normalize(imagePixels, imagePixels, 1.0, 0.0, NORM_L2, CV_32F);
    imagePixels.convertTo(img, CV_32FC1);
    cout << input_tensor.DebugString()<<endl;
 
    /*-----------------------------------用网络进行测试----------------------------------------*/
    cout<<endl<<"<-------------Running the model with test_image--------------->"<<endl;
	
	
	//前向运行，输出结果一定是一个tensor的vector
	vector<tensorflow::Tensor> outputs;
	char result[5];
    for(int i=0;i<5;i++)
	{
		 
		string output_node = output_tensor_name+to_string(i);
		Status status_run = session->Run({{input_tensor_name, input_tensor}}, {output_node}, {}, &outputs);
 
		if (!status_run.ok()) {
			cout << "ERROR: RUN failed..."  << std::endl;
			cout << status_run.ToString() << "\n";
			return -1;
		}
		int class_id=GetLabelFromOutput(outputs);
		result[i]=INDEX2CH[class_id];
		outputs.clear();
	}
	cout<<"Real code is:";
	cout<<image_path<<endl;
	cout<<"Predict result is:";
	cout<<result<<endl;
    return 0;
}