Ubuntu下，利用pb格式文件，编译基于Tensorflow的C++动态库文件

1. 简述

一个已训练的神经网络模型，要经过设计网络、选择数据集、训练、调优、固化（freeze）等过程。在这之后，为了能够让我们的神经网络，更广泛地应用于各个程序中，我们要将其编译成C++的库文件（xxx.so/xxx.a）。这样我们就可以轻松的通过C++ API接口来调用它，甚至可以不用在环境中构建tensorflow环境（静态库），方便移植到各种AIOT设备（这里有可能会涉及到交叉编译）。

本文将讲述的，是在拥有一个已经freeze的tf模型（xxx.pb）的情况下，如何将其编译成一个C++库文件。

原生Tensorflow是通过bazel工具编译的，在这里我们将用CMake来管理我们的C++程序。

2. 我们开始吧

1） 从源码编译tensorflow

是的，我们需要从源码编译tensorflow，忘了你的 pip install 和 conda install 吧！从Github上把tensorflow的源码拉下来，然后用bazel进行编译。这里有详细的过程 ：

https://tensorflow.google.cn/install/source

国内直接百度出来的链接可能没法用，好像被墙了。

2）关于CMake

可能你已经对CMake有了一定的理解，但是我们还是详细讲解一下这里要用到的CMakeLists.txt文件。

i.我们要定义一下CMake的基本设置

# ----------------------------------------------------------------------------
#   Basic Configuration
# ----------------------------------------------------------------------------
# 最低cmake版本要求
CMAKE_MINIMUM_REQUIRED(VERSION 2.8)  
# 项目的名字
PROJECT(cpp_tf_demo) 
# 定义C++标准
SET(CMAKE_CXX_STANDARD 11) 
# 这两个flag是因为我的tensorflow编译过程中产生的一些bug
SET(CMAKE_CXX_FLAGS "-Wl,--no-as-needed ${CMAKE_CXX_FLAGS}") 

ii.然后我们要设定依赖项

大部分C++工程都是依赖于已有的项目，有些项目非常完整，并且通过make install等方式安装在了计算机目录中，而另一些则可能不太完整，需要我们自己去添加引用。

对于第一种，比如说，在图像相关的项目中我们最常使用的Opencv，它就是一个非常完整的工程，我们可以通过FIND_PACKAGE来找到它的路径，同时设定其为必要项。

# ----------------------------------------------------------------------------
#   Find Dependencies
# ----------------------------------------------------------------------------
#Opencv
#寻找Opencv的引用包位置，并且设定其为必要项
FIND_PACKAGE(OpenCV  REQUIRED) 

#检测它的版本是否不低于3.0
IF(NOT OpenCV_VERSION VERSION_LESS "3.0")
    ADD_DEFINITIONS(-DOPENCV_VERSION_3)
    SET(OPENCV_VERSION_3 ON)
ELSE()
    SET(OPENCV_VERSION_3 OFF)
ENDIF()

#添加${OpenCV_INCLUDE_DIRS}到头文件夹中，如果没有这一行我们就在写程序 #include xxxx时就会报错找不到该头文件
INCLUDE_DIRECTORIES(${OpenCV_INCLUDE_DIRS})

#将${OpenCV_LIBS}添加到REQUIRED_LIBRARIES中，我们就可以正常使用Opencv中定义的模块了
SET(REQUIRED_LIBRARIES ${REQUIRED_LIBRARIES} ${OpenCV_LIBS})

对于第二种就比较复杂，比如说我们从源编译的tensorflow库，我不确定它是否可以在编译完成后，像opencv一样安装到PC目录中，所以手动添加引用命令。

# Tensorflow
# 首先我们设置一个tensorflow的主目录，就是你用git clone下来的tf代码目录，这样后续会比较方便
SET(TENSORFLOW_DIR /your/path/to/tensorflow/)
# include path
# 下面的path是我参考了一位博主的设置，然后根据自己的具体情况修改的
INCLUDE_DIRECTORIES(${TENSORFLOW_DIR})
INCLUDE_DIRECTORIES(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/gen/proto)
INCLUDE_DIRECTORIES(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/gen/protobuf-host/include)
INCLUDE_DIRECTORIES(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/downloads/eigen)
INCLUDE_DIRECTORIES(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/downloads/absl)
INCLUDE_DIRECTORIES(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/downloads/nsync/public)
# lib path
# library文件就比较简单了，直接把bazel-bin下面的tensorflow文件夹链进来就行
LINK_DIRECTORIES(${TENSORFLOW_DIR}/bazel-bin/tensorflow)

剩下的就是关于你自己编写的程序部分，我建议你创建一个inlcude文件夹和一个src文件夹，把所有需要的头文件和源文件都分别放进去（仅针对小型项目，适用于跟我一样刚开始接触这一块儿的小白，一个项目就输出一个库文件），然后在CMakeLists.txt（也就是写上面代码的文件）里添加

# Project Head Files
INCLUDE_DIRECTORIES(include)
# Project Source Files
AUX_SOURCE_DIRECTORY(src DIR_SRCS)

iii.设置你的输出库文件选项

#------------------------------------------------
# LIB
#------------------------------------------------
# 首先我们设置一下输出的库文件将在哪个目录下生成
SET(LIBRARY_OUTPUT_PATH "${PROJECT_BINARY_DIR}/lib")
# 然后设置我们输出的库文件名以及库文件包含的源文件地址，这里的SHARED是指我们将输出一个.so为后缀的动态库
ADD_LIBRARY(your-lib-name SHARED ${DIR_SRCS})
# 这一行代码的意思是，我们的库将依赖于Opencv以及tensorflow，这两个tensorflow的库文件应该是在你完成tf编译后出现在bazel-bin/tensorflow目录下的
TARGET_LINK_LIBRARIES(your-lib-name ${OpenCV_LIBS} tensorflow_cc tensorflow_framework)

#------------------------------------------------
# DIRS
#------------------------------------------------
# 我还自己写了一些小的demo，用于调试库文件代码的正确性
ADD_SUBDIRECTORY(demo)

最后我们的输出库文件将以"lib+your-lib-name+.so"的形式出现（e.g. libclassification.so）。我在demo文件夹中写了一些用于调试的代码，就简单把demo文件夹下的CMakeList.txt（跟上面主目录下的CMakeList.txt不是一个文件哦）也贴出来吧，不加以赘述了。

# LINK_LIBRARIES(${PROJECT_NAME})

SET(EXECUTABLE_OUTPUT_PATH ${PROJECT_BINARY_DIR}/bin)

ADD_EXECUTABLE(demo demo.cpp)
TARGET_LINK_LIBRARIES(demo tensorflow_cc tensorflow_framework)

ADD_EXECUTABLE(demo1 demo1.cpp)
TARGET_LINK_LIBRARIES(demo1 tensorflow_cc tensorflow_framework ${OpenCV_LIBS})

ADD_EXECUTABLE(demo2 demo2.cpp)
TARGET_LINK_LIBRARIES(demo2 ${OpenCV_LIBS} ${PROJECT_BINARY_DIR}/lib/lib+your-lib-name+.so)

3）pb文件

在主目录下，创建一个networks文件夹，将你的.pb都放在里面，方便使用。

4）主程序

如果我们极度精简，以classification为例，主要的代码就两块：classification.h 和 classification.cpp。

i. classification.h

在这里声明你类

#ifndef CLASSIFICATION_FILE_HEADER_INC
#define CLASSIFICATION_FILE_HEADER_INC

#include <fstream>
#include <utility>
#include <iostream>
#include <string>
#include <vector>
#include <tensorflow/core/platform/env.h>
#include <tensorflow/core/public/session.h>
#include "opencv2/opencv.hpp"

using namespace std;
using namespace cv;
using namespace tensorflow;
 


class Classification{
public:
    string model_path;
    int input_height;
    int input_width;
    vector<string> input_tensor_name;
    vector<string> output_tensor_name;

    Session* session;
    GraphDef graphdef;


    Classification();
    Classification(string model_path, int input_height,
                int input_width, vector<string> input_tensor_name,
                vector<string> output_tensor_name);
    ~Classification();
    void CVMat_to_Tensor(Mat img,Tensor* output_tensor,int input_rows,int input_cols);
    int Initialization();
    Tensor predict(Mat image);
};

#endif

ii. classification.cpp

编写你的主程序

#include "../include/classification.h"

Classification::Classification(){
	model_path="your/path/to/pbfile/classification.pb";
    input_height = 480;
    input_width = 640;
    // 设定pb文件中输入输出的node，输出node可以是多个，输入应该也是多个（只不过某些嵌入式开发板可能不支持，这都是后话了）
    input_tensor_name={"your-input-node-name"};
    output_tensor_name={ "your-output-node-name1",
                        "your-output-node-name2"
                        }
                    };
    Initialization();
}

Classification::Classification(string model_path,  int input_height,
							int input_width, vector<string> input_tensor_name,
							vector<string> output_tensor_name)
{
	model_path = model_path;
    input_height = input_height;
    input_width = input_width;
    input_tensor_name = input_tensor_name;
    output_tensor_name = output_tensor_name;

    Initialization();
}

Classification::~Classification(){
    session->Close();
}

void CVMat_to_Tensor(Mat img,Tensor* output_tensor,int input_rows,int input_cols)
{
    //imshow("input image",img);
    //图像进行resize处理
    resize(img,img,cv::Size(input_cols,input_rows));
    //imshow("resized image",img);

    //归一化
    img.convertTo(img,CV_32FC1);
    img=1-img/255;

    //创建一个指向tensor的内容的指针
    float *p = output_tensor->flat<float>().data();

    //创建一个Mat，与tensor的指针绑定,改变这个Mat的值，就相当于改变tensor的值
    Mat tempMat(input_rows, input_cols, CV_32FC1, p);
    img.convertTo(tempMat,CV_32FC1);
}


int Classification::Initialization(){
    Status status = NewSession(SessionOptions(), &this->session); 
    Status status_load = ReadBinaryProto(Env::Default(), model_path, &this->graphdef); //从pb文件中读取图模型;
    if (!status_load.ok()) {
        cout << "ERROR: Loading model failed..." << model_path << std::endl;
        cout << status_load.ToString() << "\n";
        return -1;
    }
    Status status_create = session->Create(this->graphdef); //将模型导入会话Session中;
    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;
    return 0;
}



Tensor Classification::predict(Mat image){
    //创建一个tensor作为输入网络的接口
 	Tensor resized_tensor(DT_FLOAT, TensorShape({1,input_height,input_width,1}));
    //将Opencv的Mat格式的图片存入tensor
    CVMat_to_Tensor(image,&resized_tensor,input_height,input_width);

    //前向运行，输出结果一定是一个tensor的vector
    vector<tensorflow::Tensor> outputs;
    Status status_run = session->Run({{input_tensor_name[0], resized_tensor}}, 
    									output_tensor_name, {}, &outputs);

    if (!status_run.ok()) {
        cout << "ERROR: RUN failed..."  << std::endl;
        cout << status_run.ToString() << "\n";
        return -1;
    }

    Tensor output = outputs[0];
    return output;
}

注1：这里的推理仅限于batchsize=1的情况下，如果你有别的需求，请自行修改代码。
注2：这不完全是我自己运行的代码，因为自己的代码和公司项目相关，我把内部相关的部分用一个简单的classification替代了，predict输出的格式也更换了，有可能存在部分小bug，需要你去自己调试了，我这里就是打个样。

5）编译你的库文件

我们在这里推荐外部编译的做法，创建一个build文件夹，所有编译生成的文件都将在出现在这个文件夹下，不会和原始文件混杂在一起。

mkdir build

进入build文件夹

cd build

利用cmake编写makefile文件

cmake  ..

开始编译

make

如果一切顺利，你将会在build/lib文件夹下看到你的.so库。

撒花！

3.后记

这里我们只介绍了如何编译动态库，而静态库编译的方法被省略了（主要原因是我也没去弄过），主要的流程应该差不多，只需省略SHARED参数。

没有办法把我自己编译通过的完整代码po出来我也很难过，朋友们出现了问题我们可以在评论中讨论。（应该不会又太难搞的bug，主体部分我没改呀！！）

中间的代码inspired by很多博主，我一时半会找不到之前看的文章了，很抱歉没办法精确地贴出引用，内心非常感谢他们的分享。