1. 为什么要在C++里加载PB模型
把训练好的模型塞进C++工程,是很多做算法落地的人绕不开的一步。Python那边调model.predict()两行代码搞定的事,到了C++这边就变成了一堆API、指针和内存管理的组合拳。我最早接触这个需求是在一个工业质检的模拟项目里,产线上的相机每秒出图,检测节拍要求单帧推理控制在30毫秒以内,Python的解释器开销和GIL锁直接把方案毙掉了,最后只能走C++加载PB模型这条路。
PB模型,全称是Protocol Buffer格式的模型文件,TensorFlow训练完之后用freeze_graph或者convert_variables_to_constants把计算图和权重固化到一个.pb文件里,这个文件就是C++端要加载的东西。它和checkpoint的区别在于,checkpoint是变量分开存的,加载时需要重新构建图结构;而PB文件把图结构和参数值打包在一起,C++端只需要一个ReadBinaryProto加Session::Create就能把整个模型跑起来,省去了在C++里重新搭图的痛苦。
这篇文章适合谁看?如果你手头有一个训练好的TensorFlow模型,想把它集成到C++的服务、桌面软件或者嵌入式程序里,又不想依赖Python运行时,那这篇内容基本能覆盖你从环境配置到推理调用的全流程。我会把踩过的坑、参数怎么选、内存怎么管这些细节都摊开讲,代码可以直接抄。
需要提前说明的是,TensorFlow的C++ API在不同版本之间差异不小,尤其是1.x和2.x的接口变化很大。我这里以TensorFlow 1.15这个在工业界仍然大量使用的版本为主线,因为它的C++ API相对稳定,资料也最全。2.x的思路类似,但部分头文件和命名空间有调整,我会在关键位置标注差异。
2. 环境准备与编译方案选型
2.1 三种集成方式的取舍
在C++里用TensorFlow,摆在面前的路其实有三条,每条路的成本和适用场景完全不同,选错了后面会非常难受。
第一条是源码编译TensorFlow C++库。从TensorFlow的源码用Bazel编译出libtensorflow_cc.so和libtensorflow_framework.so,然后链接到自己的工程里。这条路最灵活,可以自定义编译选项,但编译时间动辄几个小时,Bazel的依赖下载在国内网络环境下也是个折磨,而且编译出来的库体积巨大,动辄几百MB。
第二条是使用官方预编译的C库。TensorFlow官方提供了libtensorflow.so的预编译包,配合C API(c_api.h)使用。这个方案最轻量,下载解压就能用,但C API的功能比C++ API少很多,比如不支持直接操作Tensor对象的高级方法,很多操作要手动管理TF_Tensor的生命周期,写起来比较繁琐。
第三条是用C++ API但链接预编译的C库。这是很多人不知道的一个折中方案——TensorFlow的C++头文件(tensorflow/cc/...)其实可以配合预编译的libtensorflow.so使用,只要头文件版本和库版本一致。这样既能用C++的Session、GraphDef这些好用的类,又省去了编译源码的麻烦。我实测下来,TensorFlow 1.15的C++头文件配合官方预编译的C库,加载PB模型跑推理是完全可行的。
注意:第三条路有个前提,你的工程里不能用到C库没有导出的符号。如果只是加载PB模型做前向推理,这个方案足够;但如果要做自定义Op或者训练相关的操作,还是得走源码编译。
2.2 头文件和库的获取
以TensorFlow 1.15为例,你需要准备这些东西:
- 头文件:从TensorFlow的GitHub仓库下载对应版本的源码,只需要
tensorflow/目录下的头文件,不需要编译。关键的头文件包括tensorflow/core/public/session.h、tensorflow/core/framework/tensor.h、tensorflow/core/platform/env.h、tensorflow/cc/ops/standard_ops.h等。 - 动态库:从官方发布的预编译包中获取
libtensorflow.so和libtensorflow_framework.so。这两个库有依赖关系,libtensorflow.so依赖libtensorflow_framework.so,运行时两个都要在LD_LIBRARY_PATH里。
目录结构建议这样组织:
project/ ├── include/ │ └── tensorflow/ # 从源码拷贝的头文件 ├── lib/ │ ├── libtensorflow.so │ └── libtensorflow_framework.so ├── src/ │ └── main.cpp └── CMakeLists.txt2.3 CMake配置的关键点
CMakeLists.txt的写法直接决定了你能不能编译通过,这里有几个容易翻车的地方:
cmake_minimum_required(VERSION 3.10) project(tf_inference) set(CMAKE_CXX_STANDARD 14) set(CMAKE_CXX_STANDARD_REQUIRED ON) # 头文件路径 include_directories(${CMAKE_SOURCE_DIR}/include) # 库路径 link_directories(${CMAKE_SOURCE_DIR}/lib) add_executable(tf_inference src/main.cpp) # 链接顺序很重要,libtensorflow.so 必须在 libtensorflow_framework.so 前面 target_link_libraries(tf_inference tensorflow tensorflow_framework pthread dl )第一个坑是C++标准。TensorFlow 1.15的头文件大量使用了C++11的特性,部分地方需要C++14,如果你用C++17编译可能会遇到std::string相关的兼容性问题。我建议锁死在C++14。
第二个坑是链接顺序。GNU链接器对库的顺序敏感,libtensorflow.so依赖libtensorflow_framework.so,所以前者必须写在前面,否则会出现大量undefined reference错误。
第三个坑是ABI兼容性。TensorFlow官方预编译库是用特定版本的GCC和_GLIBCXX_USE_CXX11_ABI设置编译的。如果你自己的代码用了不同的ABI设置,链接时会报符号找不到。TensorFlow 1.15的预编译库默认是_GLIBCXX_USE_CXX11_ABI=0,你需要在CMake里加上:
add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)这个宏不加,编译能过,链接必挂,而且报错信息很难看懂,全是std::__cxx11::basic_string相关的符号缺失。
3. PB模型加载的核心流程拆解
3.1 从文件到GraphDef
加载PB模型的第一步是把磁盘上的二进制文件读进内存,解析成GraphDef对象。GraphDef是TensorFlow对计算图的序列化表示,里面记录了所有的节点、边和常量值。
#include "tensorflow/core/public/session.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/graph/default_device.h" using namespace tensorflow; int main() { // 1. 创建Session Session* session; Status status = NewSession(SessionOptions(), &session); if (!status.ok()) { std::cerr << "创建Session失败: " << status.ToString() << std::endl; return -1; } // 2. 读取PB文件 GraphDef graph_def; status = ReadBinaryProto(Env::Default(), "model.pb", &graph_def); if (!status.ok()) { std::cerr << "读取PB文件失败: " << status.ToString() << std::endl; delete session; return -1; } // 3. 将GraphDef加载到Session status = session->Create(graph_def); if (!status.ok()) { std::cerr << "加载图失败: " << status.ToString() << std::endl; delete session; return -1; } // ... 后续推理 }这三步看起来简单,但每一步都有讲究。SessionOptions里可以配置线程数、GPU选项、内存增长策略等,默认配置在CPU推理场景下够用,但如果你的模型比较大,建议显式设置config.set_intra_op_parallelism_threads()和config.set_inter_op_parallelism_threads(),前者控制单个算子内部的并行度,后者控制算子之间的并行度。我一般设成CPU核心数,实测下来比默认值快15%左右。
ReadBinaryProto读的是二进制格式的PB文件,如果你手上的PB是文本格式(很少见,但确实有),要用ReadTextProto。判断方法很简单,用文本编辑器打开,如果开头是乱码就是二进制,如果是可读的文本就是文本格式。
3.2 输入输出的张量对接
模型加载完之后,下一步是搞清楚这个模型的输入输出节点叫什么名字。这是新手最容易卡住的地方——PB文件里节点名可能是一串你看不懂的字符串,比如input_1:0、conv2d_1/Relu:0这种。
获取节点名有两个办法。第一个办法是在Python端打印:
import tensorflow as tf with tf.gfile.GFile("model.pb", "rb") as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) for node in graph_def.node: print(node.name, node.op)第二个办法是在C++端遍历graph_def.node(),打印每个节点的名字和操作类型。我一般用Python端确认,因为输出更直观,而且可以顺便用Netron这个工具可视化整个图结构,一眼就能看出输入输出在哪。
拿到节点名之后,构造输入张量的代码是这样的:
// 假设输入节点名是 "input:0",形状是 [1, 224, 224, 3] Tensor input_tensor(DT_FLOAT, TensorShape({1, 224, 224, 3})); // 填充数据 auto input_flat = input_tensor.flat<float>(); for (int i = 0; i < input_flat.size(); ++i) { input_flat(i) = image_data[i]; // image_data 是你的预处理后的图像数据 } // 构造输入feed std::vector<std::pair<std::string, Tensor>> inputs; inputs.push_back({"input:0", input_tensor}); // 指定输出节点 std::vector<std::string> output_names = {"output:0"}; // 存储输出 std::vector<Tensor> outputs; // 执行推理 Status status = session->Run(inputs, output_names, {}, &outputs); if (!status.ok()) { std::cerr << "推理失败: " << status.ToString() << std::endl; }这里有几个细节值得展开说。TensorShape的维度顺序是NHWC,也就是[batch, height, width, channels],这和OpenCV读进来的图像格式(HWC)差了一个batch维度,需要手动扩维。如果你的模型输入是NCHW格式(比如从PyTorch转过来的),那就要做相应的转置。
session->Run的第三个参数是target_node_names,一般传空vector,表示要执行所有输出节点依赖的算子。第四个参数outputs是输出容器,调用完之后outputs[0]就是第一个输出节点的结果。
3.3 输出结果的解析
推理跑完之后,outputs里存的是Tensor对象,取数据的方式取决于你的模型输出类型。分类模型一般输出[1, num_classes]的浮点数组,检测模型可能输出多个张量(boxes、scores、classes)。
// 取第一个输出 Tensor& output_tensor = outputs[0]; // 打印形状 std::cout << "输出形状: " << output_tensor.shape().DebugString() << std::endl; // 取数据 auto output_flat = output_tensor.flat<float>(); int num_classes = output_flat.size(); // 找最大值(分类场景) int max_idx = 0; float max_val = output_flat(0); for (int i = 1; i < num_classes; ++i) { if (output_flat(i) > max_val) { max_val = output_flat(i); max_idx = i; } } std::cout << "预测类别: " << max_idx << ", 置信度: " << max_val << std::endl;如果输出是DT_INT32或者其他类型,把flat<float>()换成对应的类型即可。TensorFlow的Tensor类支持flat<T>()、matrix<T>()、tensor<T, N>()等多种访问方式,选最贴合你数据维度的那种。
提示:
flat<T>()返回的是一个Eigen的Map对象,不是拷贝,直接操作它就是在操作Tensor的底层内存。如果你需要把数据传出去,记得做一次拷贝,否则Tensor析构之后指针就悬空了。
4. 完整实操:从零跑通一个图像分类模型
4.1 模型准备与冻结
假设你已经在Python端训练好了一个图像分类模型,保存成了checkpoint。第一步是把它冻结成PB文件:
import tensorflow as tf from tensorflow.python.framework import graph_util # 加载模型结构 saver = tf.train.import_meta_graph("model/model.ckpt.meta") with tf.Session() as sess: # 恢复权重 saver.restore(sess, "model/model.ckpt") # 获取输出节点名 output_node_name = "output/predictions:0" # 冻结图 output_graph_def = graph_util.convert_variables_to_constants( sess, sess.graph.as_graph_def(), [output_node_name.replace(":0", "")] ) # 写入PB文件 with tf.gfile.GFile("frozen_model.pb", "wb") as f: f.write(output_graph_def.SerializeToString()) print("PB文件已生成,大小: %.2f MB" % (len(output_graph_def.SerializeToString()) / 1024 / 1024))convert_variables_to_constants的第二个参数是要保留的输出节点名列表,注意这里要去掉:0后缀,只保留节点名。这个函数会把图中所有变量替换成常量,这样C++端加载时就不需要再初始化变量了。
冻结完之后,强烈建议用Python端先验证一下PB文件能不能正常加载和推理:
import tensorflow as tf import numpy as np with tf.gfile.GFile("frozen_model.pb", "rb") as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) with tf.Graph().as_default() as graph: tf.import_graph_def(graph_def, name="") with tf.Session() as sess: input_tensor = graph.get_tensor_by_name("input:0") output_tensor = graph.get_tensor_by_name("output/predictions:0") # 构造随机输入测试 test_input = np.random.rand(1, 224, 224, 3).astype(np.float32) result = sess.run(output_tensor, feed_dict={input_tensor: test_input}) print("Python端推理结果形状:", result.shape)这一步能过,说明PB文件本身没问题,后面C++端出问题就只可能是环境或代码的问题,排查范围大大缩小。
4.2 C++端完整推理代码
下面是一个完整的、可以直接编译运行的C++推理代码,包含了图像预处理、推理和后处理:
#include <iostream> #include <vector> #include <chrono> #include <opencv2/opencv.hpp> #include "tensorflow/core/public/session.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/framework/tensor.h" using namespace tensorflow; class ImageClassifier { public: ImageClassifier(const std::string& model_path, const std::string& input_name, const std::string& output_name, int input_height, int input_width) : input_name_(input_name), output_name_(output_name), input_height_(input_height), input_width_(input_width) { // 配置Session SessionOptions options; options.config.set_intra_op_parallelism_threads(4); options.config.set_inter_op_parallelism_threads(4); Status status = NewSession(options, &session_); if (!status.ok()) { throw std::runtime_error("创建Session失败: " + status.ToString()); } // 加载图 GraphDef graph_def; status = ReadBinaryProto(Env::Default(), model_path, &graph_def); if (!status.ok()) { throw std::runtime_error("读取PB失败: " + status.ToString()); } status = session_->Create(graph_def); if (!status.ok()) { throw std::runtime_error("加载图失败: " + status.ToString()); } std::cout << "模型加载成功: " << model_path << std::endl; } ~ImageClassifier() { if (session_) { session_->Close(); delete session_; } } std::vector<float> predict(const cv::Mat& image) { // 1. 预处理:resize + 归一化 cv::Mat resized; cv::resize(image, resized, cv::Size(input_width_, input_height_)); cv::Mat float_img; resized.convertTo(float_img, CV_32FC3, 1.0 / 255.0); // 2. 构造输入Tensor Tensor input_tensor(DT_FLOAT, TensorShape({1, input_height_, input_width_, 3})); auto input_flat = input_tensor.flat<float>(); // OpenCV是BGR,模型一般要RGB for (int h = 0; h < input_height_; ++h) { for (int w = 0; w < input_width_; ++w) { cv::Vec3f pixel = float_img.at<cv::Vec3f>(h, w); int idx = (h * input_width_ + w) * 3; input_flat(idx + 0) = pixel[2]; // R input_flat(idx + 1) = pixel[1]; // G input_flat(idx + 2) = pixel[0]; // B } } // 3. 推理 std::vector<std::pair<std::string, Tensor>> inputs; inputs.push_back({input_name_, input_tensor}); std::vector<Tensor> outputs; Status status = session_->Run(inputs, {output_name_}, {}, &outputs); if (!status.ok()) { throw std::runtime_error("推理失败: " + status.ToString()); } // 4. 取结果 auto output_flat = outputs[0].flat<float>(); std::vector<float> result(output_flat.data(), output_flat.data() + output_flat.size()); return result; } private: Session* session_ = nullptr; std::string input_name_; std::string output_name_; int input_height_; int input_width_; }; int main() { try { ImageClassifier classifier("frozen_model.pb", "input:0", "output/predictions:0", 224, 224); cv::Mat image = cv::imread("test.jpg"); if (image.empty()) { std::cerr << "图像读取失败" << std::endl; return -1; } // 计时 auto start = std::chrono::high_resolution_clock::now(); std::vector<float> probs = classifier.predict(image); auto end = std::chrono::high_resolution_clock::now(); auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start); std::cout << "推理耗时: " << duration.count() << " ms" << std::endl; // 找最大值 int max_idx = 0; float max_val = probs[0]; for (size_t i = 1; i < probs.size(); ++i) { if (probs[i] > max_val) { max_val = probs[i]; max_idx = i; } } std::cout << "预测类别: " << max_idx << ", 置信度: " << max_val << std::endl; } catch (const std::exception& e) { std::cerr << "异常: " << e.what() << std::endl; return -1; } return 0; }这段代码里有几个我特意处理的地方。第一,把Session和模型加载封装在构造函数里,析构函数里释放,避免每次推理都重新加载模型——重新加载一次PB文件在CPU上要几百毫秒,如果每帧都加载,性能直接崩掉。第二,预处理里的BGR到RGB转换是逐像素做的,虽然慢一点但逻辑清晰,如果追求极致性能可以用cv::cvtColor加指针操作优化。第三,用std::chrono做了耗时统计,方便你评估是否满足业务节拍要求。
4.3 编译与运行
编译命令:
mkdir build && cd build cmake .. make -j4运行前确保动态库路径正确:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/../lib ./tf_inference如果报error while loading shared libraries: libtensorflow_framework.so,说明LD_LIBRARY_PATH没设对,或者两个so文件的版本不匹配。检查方法是用ldd tf_inference看依赖关系,缺哪个补哪个。
5. 性能优化与内存管理实战
5.1 推理速度的瓶颈在哪
C++加载PB模型跑推理,速度瓶颈通常不在TensorFlow本身,而在数据搬运和预处理上。我做过一组对比测试,在一个224x224输入、ResNet50结构的模型上:
| 环节 | 耗时(ms) | 占比 |
|---|---|---|
| 图像解码(imread) | 8-15 | 20% |
| 预处理(resize+归一化) | 5-10 | 15% |
| Tensor构造与拷贝 | 3-5 | 10% |
| 模型前向推理 | 20-35 | 50% |
| 后处理 | 1-2 | 5% |
可以看到,前向推理只占一半左右,剩下的时间都花在了数据准备上。优化思路很明确:把图像解码和预处理放到独立线程里,和推理线程做成流水线。我用一个双缓冲队列实现过,整体吞吐量提升了将近一倍。
另一个优化点是输入Tensor的复用。每次推理都新建一个Tensor对象,会触发内存分配和释放。如果输入形状固定,可以预先分配好一个Tensor,每次推理只更新里面的数据:
// 预分配 Tensor input_tensor(DT_FLOAT, TensorShape({1, 224, 224, 3})); // 每次推理时复用 auto input_flat = input_tensor.flat<float>(); // ... 填充新数据 // 直接feed同一个Tensor这个改动看起来小,但在高频推理场景下能省下不少内存分配的开销。
5.2 内存泄漏的排查
TensorFlow C++ API里最容易出内存问题的地方是Session和Tensor的生命周期管理。NewSession出来的指针必须手动delete,session->Run产生的输出Tensor在vector析构时会自动释放,但如果你把flat<float>()的指针存下来长期使用,Tensor析构后那个指针就是野指针。
我遇到过一次诡异的内存泄漏,程序跑几个小时内存就涨到几个GB。排查下来发现是在循环里反复调用session->Create(graph_def),每次Create都会在Session内部注册新的图,旧的图不会被自动释放。正确做法是Session只Create一次,后续所有推理复用同一个Session。
注意:如果你的程序需要动态切换模型,不要反复Create,而是创建多个Session,每个Session对应一个模型,用完统一释放。
5.3 多线程推理的注意事项
TensorFlow的Session本身是线程安全的,多个线程可以同时调用同一个Session的Run方法。但有两个限制:第一,Run方法内部会加锁,高并发下可能成为瓶颈;第二,如果多个线程feed同一个Tensor对象,会有数据竞争问题。
我的做法是每个线程持有自己的输入Tensor,Session共享。如果并发量特别大,可以创建Session池,每个线程从池里取一个Session用,用完还回去。Session的创建成本不低(加载图要几百毫秒),所以池子要在程序启动时预热好。
6. 常见问题与排查速查表
6.1 编译链接类问题
| 报错信息 | 原因 | 解决方法 |
|---|---|---|
undefined reference to tensorflow::... | 链接顺序错误或库路径不对 | 确保libtensorflow.so在libtensorflow_framework.so前面,检查link_directories |
std::__cxx11::basic_string...未定义 | ABI不兼容 | 添加-D_GLIBCXX_USE_CXX11_ABI=0重新编译 |
fatal error: tensorflow/core/public/session.h: No such file | 头文件路径没配 | 检查include_directories是否指向了包含tensorflow/目录的上级目录 |
cannot find -ltensorflow | 库文件名不对 | 确认so文件名是libtensorflow.so而不是libtensorflow.so.1.15,后者需要建软链接 |
6.2 运行时类问题
| 报错信息 | 原因 | 解决方法 |
|---|---|---|
Not found: FeedInputs: unable to find feed output input:0 | 输入节点名写错 | 用Python端打印所有节点名,确认输入节点的准确名称 |
Shape mismatch | 输入Tensor形状和模型期望不符 | 检查TensorShape的维度顺序和数值,NHWC还是NCHW |
Session was not created with a graph before Run() | 忘记调用session->Create() | 在Run之前必须先Create |
error while loading shared libraries | 动态库路径不对 | 设置LD_LIBRARY_PATH或用patchelf修改rpath |
| 推理结果全是NaN | 输入数据没归一化或格式错误 | 检查预处理逻辑,确认像素值范围和通道顺序 |
6.3 几个我踩过的坑
第一个坑是PB文件的节点名带命名空间。如果你在Python端用tf.import_graph_def时指定了name="prefix",那所有节点名都会带上prefix/前缀。C++端feed的时候必须用完整的名字,少一个字符都不行。我建议在Python端冻结模型时不要加name参数,保持节点名干净。
第二个坑是BatchNorm层的is_training参数。训练时BatchNorm用的是batch统计量,推理时应该用移动平均。如果冻结模型时没有正确设置is_training=False,推理结果会完全不对。检查方法是在Python端用单张图推理,和训练时的验证结果对比,如果差异很大就是这个原因。
第三个坑是输入图像的通道顺序。OpenCV读进来是BGR,但绝大多数TensorFlow模型训练时用的是RGB。这个错误很隐蔽,因为模型不会报错,只是精度会下降几个百分点。养成习惯,预处理完打印一下第一个像素的RGB值,和Python端对比确认。
第四个坑是动态库版本冲突。如果你的系统里已经装了TensorFlow的Python包,它可能带了一套自己的so文件。C++程序运行时如果LD_LIBRARY_PATH里混入了Python包的路径,会加载到错误的版本。排查方法是用ldd看实际加载的是哪个路径下的so。
6.4 精度对齐的验证方法
C++端推理结果和Python端不一致,是最让人头疼的问题。我的验证流程是这样的:
第一步,在Python端固定随机种子,生成一组确定的输入数据,保存成npy文件。第二步,Python端推理,保存输出结果。第三步,C++端读取同一组输入数据,推理,保存输出。第四步,用Python脚本对比两组输出的最大绝对误差和相对误差。
如果误差在1e-5以内,说明对齐了;如果误差很大,按这个顺序排查:输入数据是否一致(打印前10个值对比)、预处理是否一致(归一化系数、通道顺序)、模型是否一致(MD5校验PB文件)、计算精度是否一致(float32 vs float64)。
我遇到过一次误差在1e-3左右的情况,排查了半天发现是Python端用了tf.image.resize_bilinear,而C++端用的是OpenCV的cv::resize,两者的插值算法有细微差异。换成同样的插值方法后误差就降到1e-6了。
7. 工程化落地的几点经验
把PB模型集成到C++工程里,能跑通只是第一步,真正上线还要考虑很多工程问题。
模型版本管理。生产环境经常需要更新模型,如果每次更新都重新编译程序,运维成本太高。我的做法是把模型路径做成配置项,程序启动时从配置文件读取。更新模型时只需要替换PB文件,重启服务即可。更优雅的方案是支持热加载,用一个文件监视器监听PB文件的变化,检测到更新后创建新的Session替换旧的,做到不中断服务。
异常处理。TensorFlow的C++ API用Status对象传递错误,而不是抛异常。每个可能失败的调用都要检查status.ok(),否则出错时程序会静默失败或者崩溃。我见过有人不检查ReadBinaryProto的返回值,PB文件路径写错了程序也不报错,只是推理结果全是垃圾值,排查起来很费时间。
日志记录。TensorFlow内部有自己的日志系统,默认会往stderr输出大量INFO级别的日志,在生产环境很干扰。可以在SessionOptions里设置options.config.set_log_device_placement(false),或者通过环境变量TF_CPP_MIN_LOG_LEVEL=2把日志级别调到WARNING以上。
资源限制。TensorFlow默认会尝试占用所有可用的CPU核心和内存。在容器化部署的场景下,这会导致资源争抢。通过SessionOptions的config可以限制线程数和内存使用:
SessionOptions options; options.config.set_intra_op_parallelism_threads(2); options.config.set_inter_op_parallelism_threads(2); options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.5);跨平台兼容。如果你的程序要同时在Linux和Windows上跑,注意PB文件的字节序问题。TensorFlow的PB格式本身是跨平台的,但如果你在读取PB之前做了自定义的序列化操作,要确保字节序一致。另外Windows下动态库是.dll格式,链接方式和Linux不同,CMake里要用if(WIN32)做条件判断。
我个人在实际项目中的体会是,C++加载PB模型这件事,技术难度集中在环境配置和精度对齐上,真正写推理代码的部分反而不复杂。把环境搭好、把节点名确认清楚、把预处理逻辑对齐,后面就是按部就班的工程活了。建议第一次集成的时候留出足够的时间做精度验证,不要等到上线了才发现C++端和Python端结果对不上,那时候排查成本会高很多。