tensorflow的save和restore

Posted 2021-01-02 shixisheng

使用tensorflow中的save和restore可以对模型进行保存和恢复

import tensorflow as tf

v1 = tf.Variable(tf.random_normal([1,2]), name="v1")
v2 = tf.Variable(tf.random_normal([2,3]), name="v2")
init_op = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
    sess.run(init_op)
    print ("V1:",sess.run(v1))  
    print ("V2:",sess.run(v2))
    saver_path = saver.save(sess, "save/model.ckpt")
    print ("Model saved in file: ", saver_path) 

运行如下：

可以在本地目录发现，程序自己生成了文件夹save，然后里边有四个文件

然后从该目录下读取存储下来的数据，然后输出

import tensorflow as tf
v1 = tf.Variable(tf.random_normal([1,2]), name="v1")
v2 = tf.Variable(tf.random_normal([2,3]), name="v2")
saver = tf.train.Saver()

with tf.Session() as sess:
    saver.restore(sess, "save/model.ckpt")
    print ("V1:",sess.run(v1))
    print ("V2:",sess.run(v2))
    print ("Model restored")

从输出结果可以看到，输出的变量值就是之前存储下来的变量值，而且从模型中恢复参数时，就不用写，参数初始化语句了，

init_op = tf.global_variables_initializer()，通过从本地读取出来，就可以完成对变量的赋值。