使用tf.image.Dataset时在TensorFlow中冻结计算图的技术问询

嘿，我整理了一段封装好的函数，用来借助tensorflow.python.tools.freeze_graph工具完成TensorFlow计算图的冻结与保存，把训练后的模型结构和权重整合到一个文件里，方便后续部署使用：

import os
import tensorflow as tf
from tensorflow.python.tools import freeze_graph

def freeze_and_save_graph(self, session, save_dir, name):
    # 定义各类文件的路径与名称
    checkpoint_prefix = os.path.join(save_dir, "model")
    checkpoint_state_name = "checkpoint"
    input_graph_name = "input_graph.pbtxt"
    output_graph_name = name
    
    # 第一步：保存模型权重到检查点文件
    checkpoint_path = self.saver.save(
        session,
        checkpoint_prefix,
        global_step=0,
        latest_filename=checkpoint_state_name
    )
    
    # 第二步：导出未冻结的计算图结构到文本格式文件
    tf.train.write_graph(session.graph_def, save_dir, input_graph_name, as_text=True)
    
    # 第三步：调用freeze_graph工具完成计算图冻结
    freeze_graph.freeze_graph(
        input_graph=os.path.join(save_dir, input_graph_name),
        input_saver="",
        input_binary=False,
        input_checkpoint=checkpoint_path,
        # 重点！替换成你自己模型的输出节点名称，比如"predictions"或"output"
        output_node_names="your_model_output_node_name",
        restore_op_name="save/restore_all",
        filename_tensor_name="save/Const:0",
        output_graph=os.path.join(save_dir, output_graph_name),
        clear_devices=True,
        initializer_nodes=""
    )

关键细节说明 #

• 检查点保存：通过self.saver.save()将当前会话中的模型权重保存为检查点文件，这是后续冻结操作的基础，确保权重能被正确加载。
• 计算图导出：tf.train.write_graph()会把当前会话的计算图结构导出为.pbtxt文本文件，让freeze_graph工具能读取到完整的图结构。
• 输出节点指定：output_node_names是最关键的参数，必须填写你模型的实际输出节点名称——工具会从这个节点回溯，把所有依赖的计算节点和权重都整合到冻结的图文件中，填错的话会导致冻结后的模型无法正常使用。
• 设备信息清理：clear_devices=True会移除计算图中绑定的特定设备信息，让冻结后的模型可以在不同的硬件设备上加载运行。

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