tf.train.Saver提供Save和Restore Tensorflow变量的功能,常用于保存、还原模型训练结果,这在自己的训练和迁移学习中都很有用。
训练、保存脚本:
import tensorflow as tf checkpoint_dir = './ckpt/' x_train = [1, 2, 3, 6, 8] y_train = [4.8, 8.5, 10.4, 21.0, 25.3] x = tf.placeholder(tf.float32, name='x') y = tf.placeholder(tf.float32, name='y') W = tf.Variable(1, dtype=tf.float32, name='W') b = tf.Variable(0, dtype=tf.float32, name='b') # 定义模型 linear_model = W * x + b
with tf.name_scope("loss-model"): loss = tf.reduce_sum(tf.square(linear_model - y)) acc = tf.sqrt(loss) train_step = tf.train.GradientDescentOptimizer(0.001).minimize(loss) sess = tf.Session() init = tf.global_variables_initializer() sess.run(init) variable_saver = tf.train.Saver(max_to_keep=3) # 训练、保存variables for i in range(1000): sess.run([train_step], {x: x_train, y: y_train}) if i%10 == 0: variable_saver.save(sess, checkpoint_dir, i) curr_W, curr_b, curr_loss, curr_acc = sess.run([W, b, loss, acc], {x: x_train, y: y_train}) print("After train W: %f, b: %f, loss: %f, acc: %f" % (curr_W, curr_b, curr_loss, curr_acc))
运行保存的文件如下
还原保存的变量:
import tensorflow as tf checkpoint_dir = './ckpt/' W = tf.Variable(1, dtype=tf.float32, name='W') b = tf.Variable(0, dtype=tf.float32, name='b') sess = tf.Session() init = tf.global_variables_initializer() sess.run(init) variable_saver = tf.train.Saver(max_to_keep=3) latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir) if latest_checkpoint is not None: variable_saver.restore(sess, latest_checkpoint) curr_W, curr_b = sess.run([W, b]) print("After train W: %f, b: %f" % (curr_W, curr_b))
参考了:https://blog.csdn.net/gzj_1101/article/details/80299610