如何在C++代码中运行Yolov5导出的TensorFlow格式model.pb模型

1. 基于TensorFlow原生C++ API调用 #

• 前置准备：编译和你导出PB模型时版本完全一致的TensorFlow C库，编译时勾选C接口编译选项，避免版本不兼容导致模型加载失败。
• 核心调用逻辑示例：

#include "tensorflow/core/public/session.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/framework/tensor.h"

using namespace tensorflow;

int main() {
  // 初始化Session
  Session* session;
  Status status = NewSession(SessionOptions(), &session);
  if (!status.ok()) {
    std::cout << status.ToString() << "\n";
    return 1;
  }

  // 读取PB冻结图文件
  GraphDef graph_def;
  status = ReadBinaryProto(Env::Default(), "你的model.pb文件路径", &graph_def);
  if (!status.ok()) {
    std::cout << status.ToString() << "\n";
    return 1;
  }

  // 加载计算图到Session
  status = session->Create(graph_def);
  if (!status.ok()) {
    std::cout << status.ToString() << "\n";
    return 1;
  }

  // 构造输入Tensor，需和导出PB时的输入维度、格式完全匹配，默认YOLOv5导出输入为1*640*640*3的NHWC格式
  Tensor input(DT_FLOAT, TensorShape({1, 640, 640, 3}));
  auto input_tensor = input.tensor<float, 4>();
  // 此处补充你的图像预处理逻辑：读取图像、resize、归一化到0-1范围，将像素值填充到input_tensor中

  // 执行推理
  std::vector<std::pair<string, Tensor>> inputs = {{"你的输入节点名", input}};
  std::vector<string> output_nodes = {"你的输出节点名"};
  std::vector<Tensor> outputs;
  status = session->Run(inputs, output_nodes, {}, &outputs);
  if (!status.ok()) {
    std::cout << status.ToString() << "\n";
    return 1;
  }

  // 后处理逻辑：解析输出Tensor，执行NMS、坐标还原、类别筛选，和Python端YOLOv5后处理逻辑完全一致即可
  auto output_tensor = outputs[0].tensor<float, 3>();

  // 释放资源
  session->Close();
  delete session;
  return 0;
}

• 输入/输出节点名可通过TensorFlow自带的saved_model_cli工具查询，无需额外查找外部资料。

2. 基于OpenCV DNN模块调用（更简易、性能更优） #

如果觉得编译TensorFlow C++库复杂度太高，可以选择OpenCV DNN模块加载PB模型，4.0以上版本的OpenCV原生支持TensorFlow冻结图格式，依赖配置难度低，CPU推理速度通常比原生ONNX Runtime高15%-30%，匹配你的速度优化需求。

• 核心调用逻辑示例：

#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>

using namespace cv;
using namespace cv::dnn;

int main() {
  // 加载PB模型
  Net net = readNetFromTensorflow("你的model.pb文件路径");
  // 配置推理后端，支持CPU/GPU加速，有N卡可替换为DNN_BACKEND_CUDA+DNN_TARGET_CUDA进一步提速
  net.setPreferableBackend(DNN_BACKEND_OPENCV);
  net.setPreferableTarget(DNN_TARGET_CPU);

  // 构造输入
  Mat img = imread("测试图像路径");
  Mat blob = blobFromImage(img, 1/255.0, Size(640, 640), Scalar(0,0,0), true, false);
  net.setInput(blob);

  // 执行推理
  Mat output = net.forward();

  // 后处理逻辑：解析输出Mat，完成检测框筛选、NMS、坐标还原即可
  return 0;
}

注意事项 #

• 导出PB时需确认导出的是冻结图（frozen graph）格式，剔除所有训练相关节点，否则会出现加载失败问题。
• 前后处理逻辑必须和Python端完全对齐，包括归一化系数、通道顺序、输入尺寸，否则会出现推理结果异常。

内容的提问来源于stack exchange，提问作者M.Akyuzlu