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})
主要有三个步骤:
- 加载保存的pb模型文件,生成
graph
和session
- 使用
OpenCV
读取图像进行处理,将OpenCV
的Mat
转成适合模型读入的tensor
类型 - 放入网络
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;
}