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

最新推荐文章于 2022-04-29 10:30:31 发布

gqixl

最新推荐文章于 2022-04-29 10:30:31 发布

阅读量250

点赞数 1

分类专栏： tensorflow

43 篇文章

订阅专栏

本文介绍如何在TensorFlow中保存和恢复模型。通过实例化tf.train.Saver，可在训练过程中保存模型，并在后续训练或测试中恢复。文章提供了一个完整的示例，演示了模型保存与恢复的具体步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

深度学习平台：TensorFlow

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

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

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

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

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

实际上每调用一次保存操作会创建后3个数据文件并创建一个检查点（checkpoint）文件，简单理解就是权重等参数被保存到 .ckpt.data 文件中，以字典的形式；图和元数据被保存到 .ckpt.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))