简介
最近在研究如何打通tensorflow线下 python 的脚本训练建模, 利用freeze_graph工具输出.pb图文件,之后再线上生产环境用C++代码直接调用预先训练好的模型完成预测的工作,而不需要用自己写的Inference的函数。因为目前tensorflow提供的C++的API比较少,所以参考了几篇已有的日志,踩了不少坑一并记录下来。写了一个简单的ANN模型对Iris数据集分类的Demo。
梳理过后的流程如下:
- 1. python脚本中定义自己的模型,训练完成后将tensorflow graph定义导出为protobuf的二进制文件或文本文件(一个仅有tensor定义但不包含权重参数的文件);
- 2. python脚本训练过程保存模型参数文件 *.ckpt。
- 3. 调用tensorflow自带的freeze_graph.py 小工具, 输入为格式*.pb或*.pbtxt的protobuf文件和*.ckpt的参数文件,输出为一个新的同时包含图定义和参数的*.pb文件;这个步骤的作用是把checkpoint .ckpt文件中的参数转化为常量const operator后和之前的tensor定义绑定在一起。
- 4. 在C++中新建Session,只需要读取一个绑定后的模型文件.pb, 进行预测,利用Session->Run()获得输出的tensor的值就okay;
- 5. 编译和运行,这时有两个选择:
- a) 一种是在tensorflow源代码的子目录下新建自己项目的目录和代码,然后用bazel来编译成一个很大的100多MB的二进制文件,这个方法的缺点在于无法把预测模块集成在自己的代码系统和编译环境如cmake, bcloud中,迁移性和实用性不强;参考: (https://medium.com/jim-fleming/loading-a-tensorflow-graph-with-the-c-api-4caaff88463f) 如果打不开貌似有中文翻译版的博客
- b) 另一种是自己把tensorflow源代码编译成一个.so文件,然后在自己的C++代码环境中依赖这个文件完成编译。C的API依赖libtensorflow.so,C++的API依赖libtensorflow_cc.so
运行成功后
下面通过具体的例子写了一个简单的ANN预测的demo,应该别的模型也可以参考或者拓展C++代码中的基类。测试环境:MacOS, 需要依赖安装:tensorflow, bazel, protobuf , eigen(一种矩阵运算的库);
配置环境
系统安装 HomeBrew, Bazel, Eigen
下载编译tensorflow源码
在等待30多分钟后, 如果编译成功,在tensorflow根目录下出现 bazel-bin, bazel-genfiles 等文件夹, 按顺序执行以下命令将对应的libtensorflow_cc.so文件和其他文件拷贝进入 /usr/local/lib/ 目录
这一步完成后,我们就准备好了libtensorflow_cc.so文件等,后面在自己的C++编译环境和代码目录下编译时链接这些库即可。
1. Python线下定义模型和训练
我们写了一个简单的脚本,来训练一个包含1个隐含层的ANN模型来对Iris数据集分类,模型每层节点数:[5, 64, 3],具体脚本参考项目:
https://github.com/rockingdingo/tensorflow-tutorial
1.1 定义Graph中输入和输出tensor名称
为了方便我们在调用C++ API时,能够准确根据Tensor的名称取出对应的结果,在python脚本训练时就要先定义好每个tensor的tensor_name。 如果tensor包含命名空间namespace的如"namespace_A/tensor_A" 需要用完整的名称。(Tips: 对于不清楚tensorname具体是什么的,可以在输出的 .pbtxt文件中找对应的定义); 这个例子中,我们定义以下3个tensor的tensorname
1.2 输出graph的定义文件*.pb和参数文件 *.ckpt
我们要在训练的脚本nn_model.py中加入两处代码:第一处是将tensorflow的graph_def保存成./models/目录下一个文件nn_model.pbtxt, 里面包含有图中各个tensor的定义名称等信息。 第二处是在训练代码中加入保存参数文件的代码,将训练好的ANN模型的权重Weight和Bias同时保存到./ckpt目录下的*.ckpt, *.meta等文件。最后执行 python nn_model.py 就可以完成训练过程
最后利用tensorflow自带的 freeze_graph.py小工具把.ckpt文件中的参数固定在graph内,输出nn_model_frozen.pb
脚本中的参数解释:
- --input_graph: 模型的图的定义文件nn_model.pb (不包含权重);
- --input_checkpoint: 模型的参数文件nn_model.ckpt;
- --output_graph: 绑定后包含参数的图模型文件 nn_model_frozen.pb;
- -- output_node_names: 输出待计算的tensor名字【重要】;
发现tensorflow不同版本下运行freeze_graph.py 脚本时可能遇到的Bug挺多的,列举一下:
最后如果输出如下: Converted variables to const ops. * ops in the final graph 就代表绑定成功了!发现绑定了参数的的.pb文件大小有10多MB。
2. C++API调用模型和编译
在C++预测阶段,我们在工程目录下引用两个tensorflow的头文件:
2.1 C++加载模型
在这个例子中我们把C++的API方法都封装在基类里面了。 FeatureAdapterBase 用来处理输入的特征,以及ModelLoaderBase提供统一的模型接口load()和predict()方法。然后可以根据自己的模型可以继承基类实现这两个方法,如本demo中的ann_model_loader.cpp。可以参考下,就不具体介绍了。
a) 新建Session, 从model_path 加载*.pb模型文件,并在Session中创建图。预测的核心代码如下:
b) 预测阶段的函数调用 session->Run(input_feature.input, {output_node}, {}, &outputs);
参数 const FeatureAdapterBase& input_feature, 内部的成员input_feature.input是一个Map型, std::vector<std::pair >, 类似于python里的feed_dict={"x":x, "y": y},这里的C++代码中的输入tensor_name也一定要和python训练脚本中的一致, 如上文中设定的"inputs", "targets" 等。调用基类 FeatureAdapterBase中的方法assign(std::string, std::string tname, std::vector* vec) 函数来定义。
参数 const std::string output_node, 对应的就是在python脚本中定义的输出节点的名称,如"name_scope/output_node"
记得我们之前预先编译好的libtensorflow_cc.so文件,要成功编译需要链接那个库。 运行下列命令:
参数含义:
- a) -I/usr/local/include/tf # 依赖的include文件
- b) -L/usr/local/lib/libtensorflow_cc # 编译好的libtensorflow_cc.so文件所在的目录
- c) -ltensorflow_cc # .so文件的文件名
为了方便调用,尝试着写了一个Makefile文件,将里面的路径换成自己的,每次直接用make命令执行就好
此外,在直接用g++来编译的过程中可能会遇到一些Bug, 现在记录下来
3. 运行
最后试着运行一下之前编译好的可执行文件 tfcpp_demo
我们试着预测一个样本[1,1,1,1,1],输出该样本对应的分类和概率。进行到这一步,我们终于成功完成了在python中定义模型和训练,然后 在C++生产代码中进行编译和调用的流程。