tensorflow保存模型为pb文件的各种方式

最新推荐文章于 2024-07-21 07:15:00 发布

wanghua609

最新推荐文章于 2024-07-21 07:15:00 发布

阅读量7.8k

点赞数 3

本文链接：https://blog.csdn.net/weixin_38145317/article/details/94637649

版权

本文介绍了TensorFlow中.pb文件的两种类型及其区别，包括仅保存计算图结构和包含变量值的FrozenGraphDef。详细阐述了如何将模型保存为.pb文件，以及在测试和恢复模型时的注意事项。此外，还提到了ckpt文件的作用和如何从ckpt恢复网络并保存为.pb。

摘要由CSDN通过智能技术生成

*.pb,官方描述如下:

GraphDef(.pb)-a protobuf that represents the Tensorflow training and or computation graph. This contains operators, tensors, and variables definitions.

FrozenGraphDef - a subclass of GraphDef that contains no variables. A GraphDef can be converted to a frozen graphdef by taking a checkpoint and a graphdef and converting every variable into a constant with the value looked up in the checkpoint.

这里可以简单理解为*.pb文件有两种情况,一种是仅保存了计算图结构,不包含变量值,可以通过如下代码生成

tf.train.write_graph()

还有一种就是上面提到的FrozenGraphDef, 不仅包含计算图结构,还包含了训练产生的变量值,这类*.pb可以直接被加载用于推理计算,

1.将一个图直接保存为pb形式，这个在工作目录下保存了一个名为pb_file_pathmodel.pb的文件

import tensorflow as tf
import os
from tensorflow.python.framework import graph_util

pb_file_path = os.getcwd()

with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='x')
    y = tf.placeholder(tf.int32, name='y')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    # 这里的输出需要加上name属性
    op = tf.add(xy, b, name='op_to_store')#最终输出x*y+b, b的值默认是1

    sess.run(tf.global_variables_initializer())

    # convert_variables_to_constants 需要指定output_node_names，list()，可以多个
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['x','op_to_store'])

    # 测试 OP
    feed_dict = {x: 10, y: 3}
    print(sess.run(op, feed_dict))
    print(pb_file_path)

    # 写入序列化的 PB 文件
    with tf.gf