Tensorflow模型持久化的代码实现

来源:互联网 发布:人机界面编程难不难 编辑:程序博客网 时间:2024/06/07 10:52

1.存储模型

import tensorflow as tfv1 = tf.Variable(tf.constant(1.0, shape=[1]),name="v1")v2 = tf.Variable(tf.constant(2.0, shape=[1]),name="v2")result = v1 + v2init_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.加载模型

import tensorflow as tfv1 = tf.Variable(tf.constant(1.0, shape=[1]),name="v1")v2 = tf.Variable(tf.constant(4.0, shape=[1]),name="v2")result = v1 + v2saver = tf.train.Saver()with tf.Session() as sess:    saver.restore(sess, "test/test.ckpt")#加载已经保存的模型,并通过已经保存的模型中的变量的值来计算加法    print(sess.run(result))    #计算输出结果为[3.]不是[5.]是因为该程序加载了之前保存的模型,所以计算的变量也是保存的模型中的变量

这里写图片描述


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

    import tensorflow as tfsaver = 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.加载或者保存部分变量

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

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

# 这里声明的变量名称和已经保存的模型中变量的名称不同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默认构造的函数来加载保存的模型,那么程序会报变量找不到的错误。因为保存时候的变量和加载时的变量的名称不一致。因此可以通过字典来将保存时和加载是的变量联系起来。