我爱编程

Tensorflow的模型保存和读取tf.train.Saver

2018-05-24  本文已影响52人  听风1996

目标:训练网络后想保存训练好的模型,以及在程序中读取以保存的训练好的模型。

简介

首先,保存和恢复都需要实例化一个 tf.train.Saver。
saver = tf.train.Saver()

然后,在训练循环中,定期调用 saver.save() 方法,向文件夹中写入包含了当前模型中所有可训练变量的 checkpoint 文件。
saver.save(sess, FLAGS.train_dir, global_step=step)

之后,就可以使用 saver.restore() 方法,重载模型的参数,继续训练或用于测试数据。
saver.restore(sess, FLAGS.train_dir)
一次 saver.save() 后可以在文件夹中看到新增的四个文件,


实际上每调用一次保存操作会创建后3个数据文件并创建一个检查点(checkpoint)文件,简单理解就是权重等参数被保存到 .chkp.data 文件中,以字典的形式;图和元数据被保存到 .chkp.meta 文件中,可以被 tf.train.import_meta_graph 加载到当前默认的图。

示例

下面代码是简单的保存和读取模型:(不包括加载图数据)

import tensorflow as tf  
import numpy as np  
import os  
  
#用numpy产生数据  
x_data = np.linspace(-1,1,300)[:, np.newaxis] #转置  
noise = np.random.normal(0,0.05, x_data.shape)  
y_data = np.square(x_data)-0.5+noise  
  
#输入层  
x_ph = tf.placeholder(tf.float32, [None, 1])  
y_ph = tf.placeholder(tf.float32, [None, 1])  
  
#隐藏层  
w1 = tf.Variable(tf.random_normal([1,10]))  
b1 = tf.Variable(tf.zeros([1,10])+0.1)  
wx_plus_b1 = tf.matmul(x_ph, w1) + b1  
hidden = tf.nn.relu(wx_plus_b1)  
  
#输出层  
w2 = tf.Variable(tf.random_normal([10,1]))  
b2 = tf.Variable(tf.zeros([1,1])+0.1)  
wx_plus_b2 = tf.matmul(hidden, w2) + b2  
y = wx_plus_b2  
  
#损失  
loss = tf.reduce_mean(tf.reduce_sum(tf.square(y_ph-y),reduction_indices=[1]))  
train_op = tf.train.GradientDescentOptimizer(0.1).minimize(loss)  
  
#保存模型对象saver  
saver = tf.train.Saver()  
  
#判断模型保存路径是否存在,不存在就创建  
if not os.path.exists('tmp/'):  
    os.mkdir('tmp/')  
  
#初始化  
with tf.Session() as sess:  
    if os.path.exists('tmp/checkpoint'):         #判断模型是否存在  
        saver.restore(sess, 'tmp/model.ckpt')    #存在就从模型中恢复变量  
    else:  
        init = tf.global_variables_initializer() #不存在就初始化变量  
        sess.run(init)  
  
    for i in range(1000):  
        _,loss_value = sess.run([train_op,loss], feed_dict={x_ph:x_data, y_ph:y_data})  
        if(i%50==0):  
            save_path = saver.save(sess, 'tmp/model.ckpt')  
            print("迭代次数:%d , 训练损失:%s"%(i, loss_value))  

注:
1.saver 的操作必须在 sess 建立后进行。
2.model.ckpt 必须存在给定文件夹中,'tmp/model.ckpt' 这里至少要有一层文件夹,否则无法保存。
3.恢复模型时同保存时一样,是 ‘tmp/model.ckpt’,和那3个文件名都不一样。

上一篇下一篇

猜你喜欢

热点阅读