前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >Tensorflow模型保存和读取tf.train.Saver

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

作者头像
周小董
发布2019-03-25 10:43:50
3.4K0
发布2019-03-25 10:43:50
举报
文章被收录于专栏:python前行者

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

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

saver = tf.train.Saver()

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

saver.save(sess, save_path, global_step=step)

之后,就可以使用 saver.restore() 方法,重载模型的参数,继续训练或用于测试数据。

saver.restore(sess, save_path)

模型的恢复用的是restore()函数,它需要两个参数restore(sess, save_path),save_path指的是保存的模型路径。我们可以使用tf.train.latest_checkpoint()来自动获取最后一次保存的模型。如:

代码语言:javascript
复制
model_file=tf.train.latest_checkpoint('ckpt/')
saver.restore(sess,model_file)

一次 saver.save() 后可以在文件夹中看到新增的四个文件,

image.png
image.png

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

saver.restore()时填的文件名,因为在saver.save的时候,每个checkpoint会保存三个文件,如 my-model-10000.meta, my-model-10000.index, my-model-10000.data-00000-of-00001

import_meta_graph时填的就是meta文件名,我们知道权值都保存在my-model-10000.data-00000-of-00001这个文件中,但是如果在restore方法中填这个文件名,就会报错,应该填的是前缀,这个前缀可以使用tf.train.latest_checkpoint(checkpoint_dir)这个方法获取。

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

代码语言:javascript
复制
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))

注:

  • saver 的操作必须在 sess 建立后进行。
  • model.ckpt 必须存在给定文件夹中,‘tmp/model.ckpt’ 这里至少要有一层文件夹,否则无法保存。
  • 恢复模型时同保存时一样,是 ‘tmp/model.ckpt’,和那3个文件名都不一样。
本文参与 腾讯云自媒体同步曝光计划,分享自作者个人站点/博客。
原始发表:2019年01月03日,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 作者个人站点/博客 前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
相关产品与服务
图数据库 KonisGraph
图数据库 KonisGraph(TencentDB for KonisGraph)是一种云端图数据库服务,基于腾讯在海量图数据上的实践经验,提供一站式海量图数据存储、管理、实时查询、计算、可视化分析能力;KonisGraph 支持属性图模型和 TinkerPop Gremlin 查询语言,能够帮助用户快速完成对图数据的建模、查询和可视化分析。
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档