利用Tensorflow Python接口训练可以得到ckpt模型文件,这就是本文的起点。我们假设读者机器上已经有可用的C++ Tensorflow库。用Tensorflow的C++接口做inference的大致流程如下:
- 用python脚本将ckpt等模型文件转换生成.pb文件
- 数据准备
- 运行模型
- 必要的后处理
ckpt转为pb文件
某些时候这样的步骤并不是必须的,因为Tensorflow C++接口也是支持直接读入ckpt模型的,但是我们依然推荐这样做:
- ckpt是训练模型得到的中间结果,设计本身非常方便于模型的retrain,但对于inference而言却有不少冗余的信息,如loss函数的计算,这一部分在做inference时可以剔除。
- Tensorflow Python接口允许用户利用Tensorflow的PyFunc接口把自己写的python函数打包成Tensorflow的op保存到计算图中,这些op自然也会被保存在ckpt中,可惜,这些op在Tensorflow C++接口中并不支持,直接读入ckpt会读取失败,如果遇到诸如
Not found: Op type not registered 'PyFunc' in binary running on ...
的错误,那就说明模型中含有这样的op。
- 在将ckpt模型转为pb模型的过程中可以指定模型的输入和输出,过程用没有使用到的op会被剔除,因此只要输入->输出过程中使用的都是Tensorflow的原生op,那么生成的pb模型就是可以用C++读取的。
ckpt转pb文件用python写即可,网上教程一大堆,比如这儿,这里就不啰嗦了。
C++读取模型
转换后的模型应该包含一个.pb文件和一个variables目录,c++读取模型的时候只需要指定pb和variables所在的目录即可。读取流程如下:
/*
此处应有
#include "tensorflow/core/public/session.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
*/
// contains $model_dir/**.pb and $model_dir/variables
const std::string& graph_fn = model_dir_;
// prepare session
tensorflow::SessionOptions sess_options;
tensorflow::RunOptions run_options;
tensorflow::SavedModelBundle bundle;
TF_CHECK_OK(tensorflow::LoadSavedModel(sess_options, run_options, \
graph_fn, {tensorflow::kSavedModelTagServe}, &bundle));
tensorflow::MetaGraphDef graph_def = bundle.meta_graph_def;
std::unique_ptr<tensorflow::Session>& sess = bundle.session;
模型被读取时会自动构建一个ModelBundle,其中包含graph本身以及Session等,对于运行模型而言,这已经足够了。
数据准备
正如Python接口一样,模型的输入必须是tensor,即tensorflow::Tensor。如果你的数据是用OpenCV读取的图片,那么数据类型一定是cv::Mat,你需要一些魔法把他转成Tensor类型,比如这种魔法,当然,这样的魔法是有点蠢的,用for循环把数据一个个复制过去非常耗时,正如下文中提到的,我们在后处理时可以通过找到Tensor实际存储数据的内存地址直接读取而不必复制,我想给Tensor赋值也是有同样方法的,但时间所限,还没有研究。
运行模型
/*
此处应有
#include "tensorflow/core/framework/tensor.h"
*/
std::vector<tensorflow::Tensor> output_tensor;
TF_CHECK_OK(sess -> Run({{"Print:0", image_tensor}, }, {"strided_slice_4:0", "Shape:0"}, {}, &output_tensor));
可以说是很简单粗暴了,和Python跑模型的代码几乎一模一样:
- 第一个参数是输入,每个内层的大括号都是两个元素,第一个是op的名称,第二个是输入的数据,也就是上一步中准备好的数据
- 第二个参数是输出哪些op,数量不限,可以输出任何模型中有的op,可以不必只是最终的输出
- 第三个参数跑模型用不到,留空吧
- 第四个参数是指定输出保存到哪里,它的类型必须是Tensor的vector,你的第二个参数指定了输出几个op,最终得到的output_tensor的长度就是几。
注意
- TF_CHECK_OK是个好同志,几乎所有Tensorflow的代码都可以用它包起来,因为这些函数都会返回tensorflow::Status,如果运行有问题TF_CHECK_OK可以即时报错,因为很多函数运行错误不会报错,但是会悄咪咪的return,你会发现最终结果不对,但是哪里出了问题完全看不出来。
- ckpt模型转换为pb模型时会有一步指定op的name,但是运行模型时写的op的name,并不是你在转换模型时指定的name,而是op在ckpt图中的name,搞错了的话可能会有诸如
op not found
之类的错误。
必要的后处理
后处理是什么并不重要,重要的是tensorflow::Tensor提供的操作数据可用接口少得可怜,所以常常需要把Tensor转为我们需要的类型,此处必要的一个接口就是tensorflow::Tensor.tensor_data().data(),没记错的话默认它返回是一个unsigned char*
类型,它返回什么无所谓,总之它就是Tensor实际存储数据的内存地址,只要你确切的知道神经网络返回的数据的类型和大小,就是把它reinterpret_cast成正确的指针类型,之后的数据都是连续存储的,别问我怎么访问。
// 刚刚说了output_tensor是一个vector,这里的示例中我输出的op有两个,所以vector长度是2
auto &score_tensor = output_tensor[0];
auto &shape_tensor = output_tensor[1];
int height = score_tensor.dim_size(0);
int width = score_tensor.dim_size(1);
int channel = score_tensor.dim_size(2);
// StringPiece是引自Google ProtoBuffer中的一个数据结构,是把一段连续的内存当做字符串来对待,
// StringPiece类中有几乎所有std::string具有的接口,比如size()什么的,但你要清楚它实际上不一定是字符串,这样做只是为了方便操作
tensorflow::StringPiece lane_buffer = score_tensor.tensor_data();
std::cout << "buffer size: " << lane_buffer.size() << "\n";
const float *lane_res = reinterpret_cast<const float*>(lane_buffer.data());
//但如果你还是要问我怎么访问,我可以告诉你lane_res就是一个数组,直接下标索引就可以了,注意大小(此处是height * width * channel)不要访问越界。
return lane_res;
网上关于Tensor的数据转换内容比较少,但个人感觉既然有办法在后处理时直接得到Tensor数据内存地址,应该一定有办法用类似的方式赋值给Tensor,而不必一一拷贝。
Demo
不存在的,等有空了再整理吧...