tensorflow模型保存和加载

之前用过的保存tensorflow模型和加载模型的方法，用一个简单的例子总结一下。
1、首先计算图定义代码如下:

import tensorflow as tf

a = tf.placeholder(dtype=tf.float64,shape=[1,2],name="input_data")
b = tf.Variable(initial_value=1.0,dtype=tf.float64,name="b")
b = tf.assign(b,2)
out = tf.add(a,b,name="output_data")

saver = tf.train.Saver()
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.save(sess,save_path=".\\model\\demo")

执行之后，在当前目录生成了一个model目录，包含4个文件，如下所示：

2、加载保存的计算图
下面的代码展示了模型加载和进行计算的方法。

import tensorflow as tf
import numpy as np

with tf.Session() as sess:
    structure = tf.train.import_meta_graph(".\\model\\demo.meta")
    structure.restore(sess,tf.train.latest_checkpoint(".\\model\\"))
    input_data = np.array([[10,10]])

    out = sess.run("output_data:0",feed_dict={"input_data:0":input_data})
    print(sess.run("b:0"))
    print(out)

3、只加载模型参数
tf.train.Saver提供了一个restore()方法，恢复之前保存的变量，但是使用restore()之前必须将计算图定义一遍。用这种方法加载神经网络模型参数的时候，必须先运行网络结构定义的部分，否则会报错。如下所示：

import tensorflow as tf
import numpy as np

# 模型定义
a = tf.placeholder(dtype=tf.float64,shape=[1,2],name="input_data")
b = tf.Variable(initial_value=1.0,dtype=tf.float64,name="b")
b = tf.assign(b,2)
out = tf.add(a,b,name="output_data")
# 模型定义结束

saver = tf.train.Saver()
with tf.Session() as sess:
    model_file=tf.train.latest_checkpoint(".\\model\\")
    saver.restore(sess,".\\model\\demo")
    input_data = np.array([[10, 10]])
    print(sess.run("output_data:0", feed_dict={"input_data:0": input_data}))

如果把模型定义部分的代码注掉，就会报错 ValueError: No variables to save

tensorflow模型保存和加载

相关文章