TensorFlow的pb模型

一、TensorFlow的pb模型

在深度学习领域，TensorFlow是一个非常流行的开源框架，用于构建和训练各种复杂的神经网络模型。TensorFlow pb模型是训练完成后通常保存的一种模型格式。本篇文章将详细介绍TensorFlow pb模型的常用处理方法。

1、模型导出

TensorFlow模型训练完成后，可以通过frozen过程保存为一个最终的pb模型。这个过程会将训练好的模型权重和计算图一起保存，以便后续使用。

2、模型加载

要使用TensorFlow pb模型，首先需要将其加载到内存中。可以使用TensorFlow提供的tf.saved_model.load函数来加载模型。这个函数会将pb模型加载为一个TensorFlow的SavedModel对象，以便后续进行推理或其他操作。

3、模型推理

加载TensorFlow pb模型后，可以对其进行推理操作。推理是指使用已经训练好的模型对新的输入数据进行预测或分类。在TensorFlow中，可以使用tf.saved_model.serve.predict函数来进行推理操作。这个函数会返回模型的预测结果。

4、模型优化

如果需要对TensorFlow pb模型进行优化，可以使用TensorFlow提供的优化器对其进行优化。优化器会对模型的计算图进行分析，并尝试对其进行优化以提高推理速度或减少模型大小。在TensorFlow中，可以使用tf.compat.v1.graph_editor模块来进行模型的优化操作。

5、模型转换

如果需要将TensorFlow pb模型转换为其他框架的模型格式，可以使用TensorFlow提供的转换工具。例如，可以将TensorFlow pb模型转换为ONNX格式，以便在其他支持ONNX的框架中进行推理操作。在TensorFlow中，可以使用tf.saved_model.convert工具来进行转换操作。

二、搭建虚拟环境

https://blog.csdn.net/Emins/article/details/124967944?spm=1001.2014.3001.5501

OpenCV-GitHub-dnn：https://github.com/opencv/opencv/tree/master/samples/dnn

三、从 TensorFlow 模型生成的 .pb转换为 .pbtxt

要将一个从 TensorFlow 模型生成的 .pb（Protocol Buffer 文件，通常称为 Frozen Graph）转换为 .pbtxt（Text Format Protocol Buffer 文件），你可以使用 TensorFlow 的 tf.train.write_graph 函数或者 TensorFlow 2.x 中的 tf.io.write_graph 函数。.pbtxt 文件是以文本格式存储的模型结构，这对于阅读和调试非常有帮助。

使用 TensorFlow 1.x

如果你在使用 TensorFlow 1.x，可以使用以下代码来将 .pb 文件转换为 .pbtxt 文件：

import tensorflow as tf
 
# 加载 Frozen Graph
with tf.io.gfile.GFile('path_to_your_frozen_graph.pb', "rb") as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())
 
# 将 GraphDef 写入 pbtxt 文件
tf.io.write_graph(graph_def, ".", "output_graph.pbtxt", as_text=True)

使用 TensorFlow 2.x

在 TensorFlow 2.x 中，你可以使用以下代码：

import tensorflow as tf
 
# 加载 Frozen Graph
graph_def = tf.compat.v1.GraphDef()
with tf.io.gfile.GFile('path_to_your_frozen_graph.pb', "rb") as f:
    graph_def.ParseFromString(f.read())
 
# 将 GraphDef 写入 pbtxt 文件
tf.io.write_graph(graph_def, ".", "output_graph.pbtxt", as_text=True)

注意事项

1）路径和文件名：确保替换 'path_to_your_frozen_graph.pb' 为你的 .pb 文件的实际路径。同样，输出的 .pbtxt 文件也会被保存在当前目录下，除非你指定了不同的路径。

2）TensorFlow 版本：确保你的环境中安装了正确版本的 TensorFlow。上面的代码示例适用于 TensorFlow 1.x 和 TensorFlow 2.x。如果你使用的是 TensorFlow 2.x，确保你的代码运行在兼容模式（tf.compat.v1）下，除非你完全迁移到 TensorFlow 2 的原生 API。

四、opencv调用PB模型

opencv加载tensorflow模型需要pb文件和pbtxt文件，pbtxt是可以根据pb文件生成的。

4.1 python中调用命令生成PB模型对应的pbtxt文件。

python tf_text_graph_ssd.py --input E:\opencv-4.5.5\samples\dnn\test\frozen_inference_graph.pb  --output E:\opencv-4.5.5\samples\dnn\test\frozen_inference_graph.pbtxt  --config E:\opencv-4.5.5\samples\dnn\test\pipeline.config

4.2 opencv c++调用PB模型

要在C++中使用OpenCV加载TensorFlow的.pb模型，你需要使用OpenCV的dnn模块。以下是一个简单的例子，展示如何使用OpenCV加载和使用TensorFlow模型进行图像分类。

首先，确保你已经安装了OpenCV库，并且它包括了dnn模块。

#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <iostream>
 
int main() {
    // 初始化OpenCV的dnn模块
    cv::dnn::Net net = cv::dnn::readNetFromTensorflow("model.pb", "graph.pbtxt");
    //设置计算后台
    net.setPreferableBackend(DNN_BACKEND_OPENCV);//使用opencv dnn作为后台计算
   //设置目标设备
    net.setPreferableTarget(DNN_TARGET_OPENCL);//使用OpenCL加速

    // 读取输入图像
    cv::Mat img = cv::imread("image.jpg");
 
    // 创建一个blob从图像, 准备网络输入
    cv::Mat inputBlob = cv::dnn::blobFromImage(img, 1.0, cv::Size(224, 224), cv::Scalar(), true, false);
 
    // 设置网络输入
    net.setInput(inputBlob);
 
    // 运行前向传递，获取网络输出
    cv::Mat prob = net.forward();
 
    // 找出最大概率的索引
    cv::Point classId;
    double confidence;
    cv::minMaxLoc(prob.reshape(1, 1), 0, &confidence, 0, &classId);
 
    // 输出最大概率的类别
    std::cout << "Class ID: " << classId << " Confidence: " << confidence << std::endl;
 
    return 0;
}

在这个例子中，model.pb是TensorFlow模型的权重文件，graph.pbtxt是模型的图文件（可选，如果你的模型没有这个文件，你可以只用.pb文件）。blobFromImage函数将图像转换为网络需要的格式。net.setInput设置网络输入，net.forward执行前向传递，最后minMaxLoc找出最大概率的类别。

确保你的模型输入尺寸和blobFromImage函数中的参数匹配。如果你的模型需要其他的预处理步骤（如归一化），你需要在创建blob之前进行这些步骤。

五、TensorFlow的PB模型到ONNX

1、转换工具选择

‌1.1）推荐工具：tf2onnx‌

该工具支持将TensorFlow的GraphDef（.pb）和SavedModel格式直接转换为ONNX格式‌。

安装命令：pip install -U tf2onnx‌。

‌1.2）其他辅助工具‌

‌summarize_graph‌：用于分析PB模型的输入/输出节点名称，避免手动查看模型结构‌。

2、转换步骤

场景1：GraphDef格式（冻结PB模型）

‌2.1）获取输入/输出节点名称‌

使用summarize_graph工具分析PB文件：

bazel-bin/tensorflow/tools/graph_transforms/summarize_graph --in_graph=model.pb

输出结果中查找Inputs和Outputs字段‌。

‌2.2）执行转换命令‌

python -m tf2onnx.convert \
    --graphdef model.pb \
    --output model.onnx \
    --inputs "input_node:0" \
    --outputs "output_node:0"

参数说明：--inputs和--outputs需与模型实际节点名称一致‌。

场景2：SavedModel格式

‌直接转换‌

python -m tf2onnx.convert \
    --saved-model saved_model_dir \
    --output model.onnx \
    --opset 13

无需指定输入/输出节点，工具自动解析‌。

3、注意事项

‌3.1）版本兼容性‌

TensorFlow 1.x和2.x均支持，但需确保tf2onnx版本与TensorFlow版本匹配‌。

若转换失败，可尝试调整--opset参数（如--opset 11或更高）‌。

‌3.2）节点名称准确性‌

输入/输出节点名称错误会导致转换失败或推理异常，建议通过工具或代码验证节点名称‌。

‌3.3）模型验证‌

使用onnxruntime加载转换后的ONNX模型，执行推理测试以验证正确性‌。

import onnxruntime as ort
# 加载ONNX模型
session = ort.InferenceSession('model.onnx')
# 获取输入和输出张量的名字
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
# 运行模型
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
input_data = to_numpy(your_input_data)  # 将你的输入数据转换为numpy数组
result = session.run([output_name], {input_name: input_data})[0]  # 运行模型并获取输出

在上面的代码中，你需要将'model.onnx'替换为你的ONNX模型文件的路径，'your_input_data'替换为你的输入数据。