Tensorflow:模型变量保存

2020-08-14 15:24:31 浏览数 (1)

参考文献Tensorflow 实战 Google 深度学习框架[1]实验平台: Tensorflow1.4.0 python3.5.0

Tensorflow 常用保存模型方法

代码语言:javascript复制
import tensorflow as tf
saver = tf.train.Saver()  # 创建保存器
with tf.Session() as sess:
      saver.save(sess,"/path/model.ckpt")  #保存模型到相应ckpt文件
      saver.restore(sess,"/path/model.ckpt")  #从相应ckpt文件中恢复模型变量
  • 使用 tf.train.Saver 会保存运行 Tensorflow 程序所需要的全部信息,然而有时并不需要某些信息。比如在测试或离线预测时,只需要知道如何从神经网络的输入层经过前向传播计算得到输出层即可,而不需要类似的变量初始化,模型保存等辅助节点的信息。Tensorflow 提供了 convert_varibales_to_constants 函数,通过这个函数可以将计算图中的变量及其取值通过常量的方式保存,这样整个 Tensorflow 计算图可以统一存放在一个文件中。

将变量取值保存为 pb 文件

代码语言:javascript复制
# pb文件保存方法
import tensorflow as tf
from tensorflow.python.framework import graph_util

v1 = tf.Variable(tf.constant(1.0, shape=[1]), name="v1")
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name="v2")
result = v1   v2

init_op = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init_op)  # 初始化所有变量
    # 导出当前计算图的GraphDef部分,只需要这一部分就可以完成从输入层到输出层的计算过程
    graph_def = tf.get_default_graph().as_graph_def()
    # 将需要保存的add节点名称传入参数中,表示将所需的变量转化为常量保存下来。
    output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])
    # 将导出的模型存入文件中
    with tf.gfile.GFile("Saved_model/combined_model.pb", "wb") as f:
        f.write(output_graph_def.SerializeToString())

# 2. 加载pb文件。

from tensorflow.python.platform import gfile

with tf.Session() as sess:
    model_filename = "Saved_model/combined_model.pb"
    # 读取保存的模型文件,并将其解析成对应的GraphDef Protocol Buffer
    with gfile.FastGFile(model_filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    # 将graph_def中保存的图加载到当前图中,其中保存的时候保存的是计算节点的名称,为add
    # 但是读取时使用的是张量的名称所以是add:0
    result = tf.import_graph_def(graph_def, return_elements=["add:0"])
    print(sess.run(result))
# Converted 2 variables to const ops.
# [array([3.], dtype=float32)]

参考资料

[1]Tensorflow实战Google深度学习框架: https://github.com/caicloud/tensorflow-tutorial/tree/master/Deep_Learning_with_TensorFlow/1.4.0

0 人点赞