tf.train.Saver提供Save和Restore Tensorflow变量的功能,常用于保存、还原模型训练结果,这在自己的训练和迁移学习中都很有用。
训练、保存脚本:
import tensorflow as tfcheckpoint_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)# 训练、保存variablesfor 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 tfcheckpoint_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