如何用TensorFlow 2生成frozen_inference_graph.pb及.pbtxt文件适配OpenCV

我刚解决过类似的问题，在TensorFlow 2里确实不用再依赖TF1.x的freeze_graph.py了，有更适配TF2的流程，我给你一步步拆解：

第一步：将Keras模型保存为SavedModel格式 #

TF2推荐使用SavedModel作为模型的标准存储格式，如果你还没保存，用下面的代码：

import tensorflow as tf
from tensorflow.keras.models import load_model

# 假设你的自定义模型是model
# 如果是已训练好的模型，先加载：model = load_model("your_keras_model.h5")
model.save("my_saved_model")  # 保存为SavedModel格式，会生成一个文件夹

第二步：生成冻结的frozen_inference_graph.pb文件 #

TF2里可以通过tf.compat.v1.graph_util.convert_variables_to_constants工具来冻结图（把变量转成常量），虽然是compat模块，但这是TF2官方支持的方式，具体代码如下：

import tensorflow as tf

# 加载SavedModel
loaded_model = tf.saved_model.load("my_saved_model")
infer = loaded_model.signatures["serving_default"]

# 获取图和输出节点名
graph = tf.compat.v1.get_default_graph()
output_node_names = [output.name.split(":")[0] for output in infer.outputs]

# 冻结图
frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
    sess=tf.compat.v1.Session(graph=graph),
    input_graph_def=graph.as_graph_def(),
    output_node_names=output_node_names
)

# 保存冻结后的.pb文件
with open("frozen_inference_graph.pb", "wb") as f:
    f.write(frozen_graph_def.SerializeToString())

注意：如果你的模型有多个输出节点，要确保所有输出节点都被包含进去，不然OpenCV加载时会报错。

第三步：生成对应的.pbtxt配置文件 #

.pbtxt文件用来告诉OpenCV模型的输入输出节点信息，有两种常用方法：

方法1：通过TensorBoard日志提取 #

1. 先把SavedModel的图写入日志：

writer = tf.summary.create_file_writer("log_dir")
with writer.as_default():
    tf.summary.graph(graph.as_graph_def(), step=0)
writer.close()

2. 启动TensorBoard：tensorboard --logdir=log_dir，在Graph页面里可以查看模型的节点结构，然后手动编写pbtxt，格式大概是：

model {
  node {
    name: "input_layer"  # 你的输入节点名
    op: "Placeholder"
    attr {
      key: "dtype"
      value { type: DT_FLOAT }
    }
    attr {
      key: "shape"
      value { shape { dim { size: -1 } dim { size: 224 } dim { size: 224 } dim { size: 3 } } }
    }
  }
  node {
    name: "output_layer/Softmax"  # 你的输出节点名
    op: "Softmax"
    input: "previous_layer_name"
    attr {
      key: "T"
      value { type: DT_FLOAT }
    }
  }
  # 其他节点可以按需添加，或者只保留输入输出也能让OpenCV正常加载
}

方法2：用脚本自动生成 #

你也可以用下面的脚本直接从冻结图生成pbtxt：

import tensorflow as tf

def write_pbtxt(graph_def, output_file):
    with open(output_file, "w") as f:
        for node in graph_def.node:
            f.write(f"node {{\n")
            f.write(f"  name: \"{node.name}\"\n")
            f.write(f"  op: \"{node.op}\"\n")
            for input_name in node.input:
                f.write(f"  input: \"{input_name}\"\n")
            for attr in node.attr:
                f.write(f"  attr {{\n")
                f.write(f"    key: \"{attr}\"\n")
                f.write(f"    value: {node.attr[attr]}\n")
                f.write(f"  }}\n")
            f.write(f"}}\n")

# 加载冻结的.pb文件
with open("frozen_inference_graph.pb", "rb") as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())

# 生成pbtxt
write_pbtxt(graph_def, "model_config.pbtxt")

生成后你可以手动精简一下，只保留关键的输入输出节点和必要的属性，避免冗余。

第四步：在OpenCV中加载模型 #

最后，你可以把这两个文件的内容读成缓冲区，传给cv.dnn.readNetFromTensorflow：

import cv2

# 读取.pb文件到缓冲区
with open("frozen_inference_graph.pb", "rb") as f:
    buffer_model = f.read()

# 读取.pbtxt文件到缓冲区
with open("model_config.pbtxt", "r") as f:
    buffer_config = f.read()

# 加载模型
net = cv2.dnn.readNetFromTensorflow(buffer_model, buffer_config)

这样就能在TF2环境下生成OpenCV dnn需要的文件了，亲测有效～

内容的提问来源于stack exchange，提问作者Sam