0%

Tensorflow模型持久化的代码实现

1.存储模型

1
2
3
4
5
6
7
8
9
10
11
12
import tensorflow as tf

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()
saver = tf.train.Saver()#声明tf.train.Saver()类用于保存模型

with tf.Session() as sess:
sess.run(init_op)
saver.save(sess, "test/test.ckpt")#将模型保存于test文件夹中

运行该程序后,将在test文件中出现以下的文件

这里写图片描述


2.加载模型

1
2
3
4
5
6
7
8
9
10
11
12
13
import tensorflow as tf

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

saver = tf.train.Saver()

with tf.Session() as sess:
saver.restore(sess, "test/test.ckpt")#加载已经保存的模型,并通过已经保存的模型中的变量的值来计算加法
print(sess.run(result))

#计算输出结果为[3.]不是[5.]是因为该程序加载了之前保存的模型,所以计算的变量也是保存的模型中的变量

这里写图片描述


如果不希望重复定义变量以及运算,可以直接加载已经持久化的图

1
2
3
4
5
6
7
8
	import tensorflow as tf

saver = tf.train.import_meta_graph("test/test.ckpt.meta")
with tf.Session() as sess:
saver.restore(sess, "test/test.ckpt")
#通过张量的名称来获取张量
print(sess.run(tf.get_default_graph().get_tensor_by_name("add:0")))
#输出[3.]

该程序默认保存和加载了Tensorflow计算图中定义的全部变量。


3.加载或者保存部分变量

1
saver = tf.train.Saver([v1])#该命令只用来加载变量v1

4.加载或者保存时对变量重命名

1
2
3
4
5
6
7
8
9
10
11
12
# 这里声明的变量名称和已经保存的模型中变量的名称不同
V1=tf.Variable(tf.constant(1.0,shape=[1]),name="other-v1")
V2=tf.variable(tf.constant (2.0,shape=[1]),name="other-v2")

# 如果直接使用tf.train.Saver () 来加载模型会报变量找不到的错误。下面显示了报错信息:
# tensorflow.python.framewotk.erors.NotFoundError: Tensor name "other-v2"
# not found in checkpoint files /test/test.ckpt

# 使用一个字典(aictionary) 来重命名变量可以就可以加载原来的模型了。这个字典指定了
# 原来名称为v1的变量现在加载到变量v1中(名称为other-v1),名称为v2的变量
# 加载到变量v2 中(名称为other-v2)。
saver=tf.train.Saver({"v1"=v1,"v2":v2})

如果直接通过saver=tf.train.Saver默认构造的函数来加载保存的模型,那么程序会报变量找不到的错误。因为保存时候的变量和加载时的变量的名称不一致。因此可以通过字典来将保存时和加载是的变量联系起来。