Tensorflow中变量保存与恢复(Saver)

将训练好的模型参数保存起来,以便以后进行验证或测试,这是我们经常要做的事情。tf里面提供模型保存的是tf.train.Saver()模块。

保存变量

模型保存,先要创建一个Saver对象:如

saver=tf.train.Saver()

在创建这个Saver对象的时候,有一个参数我们经常会用到,就是 max_to_keep 参数,这个是用来设置保存模型的个数,默认为5,即 max_to_keep=5,保存最近的5个模型。如果你想每训练一代(epoch)就想保存一次模型,则可以将 max_to_keep设置为None或者0,如:

saver=tf.train.Saver(max_to_keep=0)

但是这样做除了多占用硬盘,并没有实际多大的用处,因此不推荐。

当然,如果你只想保存最后一代的模型,则只需要将max_to_keep设置为1即可,即

saver=tf.train.Saver(max_to_keep=1)

创建完saver对象后,就可以保存训练好的模型了,如:

saver.save(sess,'ckpt/mnist.ckpt',global_step=step)

saver可以使用提供的计数器自动为checkpoint文件编号。这使得在训练模型时在不同的步骤保留多个检查点。在save()方法中传递可选的global_step参数,可以对checkpoint文件进行编号

saver.save(sess, 'my-model', global_step=0) ==> filename: 'my-model-0'
...
saver.save(sess, 'my-model', global_step=1000) ==> filename: 'my-model-1000'
#——————————————————例子————————————————————————
import tensorflow as tf
# 创建两个变量
v1= tf.Variable(tf.random_normal([784, 200], stddev=0.35), name="v1")
v2= tf.Variable(tf.zeros([200]), name="v2")

# 添加用于初始化变量的节点
init_op = tf.global_variables_initializer()

# Create a saver.
saver = tf.train.Saver(tf.global_variables())

# 运行图,打开会话,每1000次保存一个模型
sess = tf.Session()
for step in range(10000):
    sess.run(init_op)
    if step % 1000 == 0:
        saver.save(sess, base_path+'my-model', global_step=step)

运行结果:

这里写图片描述

恢复变量

模型的恢复用的是restore()函数,它需要两个参数restore(sess,
save_path),save_path指的是保存的模型路径。

tf.train.Saver.restore(sess, save_path)

我们可以使用tf.train.latest_checkpoint()来自动获取最后一次保存的模型。如:

model_file=tf.train.latest_checkpoint('ckpt/')
saver.restore(sess,model_file)

标红的地方,就是与保存、恢复模型相关的代码。用一个bool型变量is_continue来控制训练和继续训练两个阶段。

整个源程序:

%%python2
import sys
import tensorflow as tf
import sys,os
print(sys.version)
print(os.getcwd())

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=False)

x = tf.placeholder(tf.float32, [None, 784])
y_=tf.placeholder(tf.int32,[None,])

dense1 = tf.layers.dense(inputs=x, 
                      units=1024, 
                      activation=tf.nn.relu,
                      kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
                      kernel_regularizer=tf.nn.l2_loss)
dense2= tf.layers.dense(inputs=dense1, 
                      units=512, 
                      activation=tf.nn.relu,
                      kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
                      kernel_regularizer=tf.nn.l2_loss)
logits= tf.layers.dense(inputs=dense2, units=10, activation=None,kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),kernel_regularizer=tf.nn.l2_loss)
loss=tf.losses.sparse_softmax_cross_entropy(labels=y_,logits=logits)
train_op=tf.train.AdamOptimizer(learning_rate=0.001).minimize(loss)
correct_prediction = tf.equal(tf.cast(tf.argmax(logits,1),tf.int32), y_)    
acc= tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

sess=tf.InteractiveSession()  
sess.run(tf.global_variables_initializer())

is_train=True/False
saver=tf.train.Saver(max_to_keep=3)

if is_train:
    max_acc=0
    f=open('test/acc.txt','w')

    for i in range(100):
        batch_xs, batch_ys = mnist.train.next_batch(100)
        sess.run(train_op, feed_dict={x: batch_xs, y_: batch_ys})
        val_loss,val_acc=sess.run([loss,acc], feed_dict={x: mnist.test.images, y_: mnist.test.labels})
        print('epoch:%d, val_loss:%f, val_acc:%f'%(i,val_loss,val_acc))
        f.write(str(i+1)+', val_acc: '+str(val_acc)+'\n')

        if val_acc>max_acc:
            max_acc=val_acc
            saver.save(sess,'test/mnist.ckpt',global_step=i+1)
    f.close()

else:
    #arg:获取最近一次保存的变量文件名称
    model_file=tf.train.latest_checkpoint('test/')
    saver.restore(sess,model_file)

    max_acc=0
    f=open('test/acc2.txt','w')
    for i in range(100):
        batch_xs, batch_ys = mnist.train.next_batch(100)
        sess.run(train_op, feed_dict={x: batch_xs, y_: batch_ys})
        val_loss,val_acc=sess.run([loss,acc], feed_dict={x: mnist.test.images, y_: mnist.test.labels})
        print('epoch:%d, val_loss:%f, val_acc:%f'%(i,val_loss,val_acc))
        f.write(str(i+1)+', val_acc: '+str(val_acc)+'\n')
        if val_acc>max_acc:
            max_acc=val_acc
            saver.save(sess,'test/mnist2.ckpt',global_step=i+1)
    f.close()
sess.close()

猜你喜欢

转载自blog.csdn.net/ddy_sweety/article/details/80459756