如何恢复单文件.model格式TensorFlow预训练模型及同格式保存方法

嘿，这个问题我碰到过好几次了！那种单个.model后缀的TensorFlow模型，本质上是把计算图结构和权重参数打包在一起的文件，在TF 1.x（≥1.2）里确实有这种保存方式，下面分恢复和保存两部分给你讲清楚：

这种文件大概率是冻结图的重命名版本（把通常的.pb后缀改成了.model），也就是把训练好的权重和计算图整合在了一起，恢复步骤如下：

• 首先读取模型文件并导入计算图：

  import tensorflow as tf
  
  # 读取.model文件内容
  with tf.gfile.GFile('your_model.model', '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='')  # name=''避免给节点加前缀
  

• 接下来找到模型的输入和输出张量：
  如果你不知道输入输出节点的名字，可以先打印所有节点名称来排查：

  for op in graph.get_operations():
      print(op.name)
  

  找到对应节点后，就可以获取张量并运行模型了：

  # 示例：假设输入节点是input:0，输出节点是output:0
  input_tensor = graph.get_tensor_by_name('input:0')
  output_tensor = graph.get_tensor_by_name('output:0')
  
  # 创建会话运行模型
  with tf.Session(graph=graph) as sess:
      # 因为权重已经嵌入在图里了，不需要额外restore
      prediction = sess.run(output_tensor, feed_dict={input_tensor: your_input_data})
  

如果是TF 2.x环境，记得把上述代码中的tf替换成tf.compat.v1，并关闭Eager Execution：

import tensorflow as tf
tf.compat.v1.disable_eager_execution()
# 后续代码和上面一致，只是把tf改成tf.compat.v1

要生成这种单文件模型，核心思路是冻结计算图——把训练过程中的变量转换成常量，和图结构合并成一个文件，步骤如下：

1. 假设你已经有训练好的模型，并且持有会话对象sess，以及知道输出节点的名称：

   import tensorflow as tf
   
   # 获取当前计算图
   graph = tf.get_default_graph()
   # 示例：假设输出节点的操作名为"output_op"（不是张量名）
   output_op_name = graph.get_tensor_by_name('output:0').op.name
   

2. 冻结图，将变量转为常量：

   frozen_graph_def = tf.graph_util.convert_variables_to_constants(
       sess,  # 训练好的会话对象
       graph.as_graph_def(),
       [output_op_name]  # 指定要保留的输出节点列表
   )
   

3. 保存为.model后缀的文件：

   with tf.gfile.GFile('my_trained_model.model', 'wb') as f:
       f.write(frozen_graph_def.SerializeToString())
   

额外注意事项 #

• 如果你用的是TF 1.2~1.15之间的版本，上面的代码可以直接运行；TF 2.x环境下同样需要用tf.compat.v1兼容模块。
• 如果你的.model文件是作者用自定义方式序列化的（比如自己写了保存逻辑），那上面的方法可能不适用，但这种情况非常少见，大部分都是冻结图的变体。

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