TensorFlow保存和载入训练模型

版权声明：本文为博主原创文章，转载请附上博文链接！ https://blog.csdn.net/weixin_38314865/article/details/84781431

保存：使用saver.save()方法保存

载入：使用saver.restore()方法载入

下面是个完整例子：

保存：

import tensorflow as tf

W = tf.Variable([[1, 1, 1], [2, 2, 2]], dtype=tf.float32, name='w')
b = tf.Variable([[0, 1, 2]], dtype=tf.float32, name='b')

saver = tf.train.Saver()
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    save_path = saver.save(sess, r"D:\test\wb")  # 将W、b保存到指定位置

载入： 

import tensorflow as tf

W = tf.Variable(tf.truncated_normal(shape=(2, 3)), dtype=tf.float32, name='w')
b = tf.Variable(tf.truncated_normal(shape=(1, 3)), dtype=tf.float32, name='b')

saver = tf.train.Saver()
with tf.Session() as sess:
    saver.restore(sess, r"D:\test\wb")  # 从指定位置加载模型
    print(sess.run(W))
    print(sess.run(b))
"""
输出：
[[1. 1. 1.]
 [2. 2. 2.]]

[[0. 1. 2.]]
"""

就算W和b定义了不同于模型的值，但是仍会输出载入模型的值，如：

import tensorflow as tf

W = tf.Variable([[0,0,0],[0,0,0]],dtype = tf.float32,name='w')
b = tf.Variable([[0,0,0]],dtype = tf.float32,name='b')

saver = tf.train.Saver()
with tf.Session() as sess:
    saver.restore(sess, r"D:\test\wb")
    print(sess.run(W))
    print(sess.run(b))
"""
输出：
[[1. 1. 1.]
 [2. 2. 2.]]

[[0. 1. 2.]]
"""

这种方法不方便的在于，在使用模型的时候，必须把模型的结构重新定义一遍，然后载入对应名字的变量的值。