在训练深度模型时，我们常常需要把模型训练的参数值保存到磁盘，防止意外情况发生导致模型数据丢失。
模型保存
模型的保存很简单，可以调用tf.train.Saver().save(sess, save_path)方法将sess会话中的参数值保存到save_path路径中。
import tensorflow as tf
# 将模型写入磁盘
v1 = tf.Variable(tf.random_normal([1,2]), name='v1')
v2 = tf.Variable(tf.random_normal([2,3]), name='v2')
init_op = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init_op)
print('v1: ', sess.run(v1))
print('v2: ', sess.run(v2))
saver_path = saver.save(sess, './data/model.ckpt')
print('Model saved in file: ', saver_path)
输出结果如下：
v1: [[-0.21382442 -0.45123124]]
v2: [[-0.26286286 1.63149405 0.94820863]
[-1.51325119 0.19088188 1.83994102]]
Model saved in file: ./data/model.ckpt
模型读取
为了避免模型从头训练，我们可以提前将模型训练的中间结果保存到磁盘。如果，有意外情况发生需要中止训练，我们可以后期加载磁盘中的参数值，然后继续训练。
import tensorflow as tf
# 从磁盘读取模型
v1 = tf.Variable(tf.random_normal([1,2]), name='v1')
v2 = tf.Variable(tf.random_normal([2,3]), name='v2')
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, './data/model.ckpt')
print('v1: ', sess.run(v1))
print('v2: ', sess.run(v2))
print('Model restored')
输出结果如下：
v1: [[-0.21382442 -0.45123124]]
v2: [[-0.26286286 1.63149405 0.94820863]
[-1.51325119 0.19088188 1.83994102]]
Model restored
通过对比可以发现，两次输出的参数结果都是一样的。
basic_cnn模型保存与读取
为了结合具体的例子，这里将上一篇文章tensorflow入门应用方法(三)——卷积网络搭建中搭建的卷积模型进行参数值保存和读取实验。
修改代码
首先，修改相关代码如下：
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('data/', one_hot=True)
trainimg = mnist.train.images
trainlabel = mnist.train.labels
testimg = mnist.test.images
testlabel = mnist.test.labels
print('MNIST loaded')
# 参数初始化
n_input = 784 # 28×28
n_output = 10
# 输入输出
x = tf.placeholder(tf.float32, [None, n_input])
y = tf.placeholder(tf.float32, [None, n_output])
keepratio = tf.placeholder(tf.float32)
Weights = {
'wc1': tf.Variable(tf.random_normal([3,3,1,64], stddev=0.1)),
# 3,3,1,64 -- f_hight, f_width, input_channel, output_channel
'wc2': tf.Variable(tf.random_normal([3,3,64,128], stddev=0.1)),
'wd1': tf.Variable(tf.random_normal([7*7*128, 1024], stddev=0.1)),
# 两次max pool以后由28×28变成7×7 卷积特征核的大小变化公式：(input_(h,w) - f_(h,w) + 2*padding)/stride + 1，
# 这里卷积对特征核大小没有影响
'wd2': tf.Variable(tf.random_normal([1024, n_output], stddev=0.1))
}
biases = {
'bc1': tf.Variable(tf.random_normal([64], stddev=0.1)),
'bc2': tf.Variable(tf.random_normal([128], stddev=0.1)),
'bd1': tf.Variable(tf.random_normal([1024], stddev=0.1)),
'bd2': tf.Variable(tf.random_normal([n_output], stddev=0.1))
}
# 定义卷积层
def conv_basic(_input, _w, _b, _keepratio):
# input
_input_r = tf.reshape(_input, shape=[-1, 28, 28, 1])
# shape=[batch_size, input_height, input_width, input_channel]
# conv layer1
_conv1 = tf.nn.conv2d(_input_r, _w['wc1'], strides=[1,1,1,1], padding='SAME')
# strides = [batch_size_stride, height_stride, width_stride, channel_stride]
_conv1 = tf.nn.relu(tf.nn.bias_add(_conv1, _b['bc1']))
_pool1 = tf.nn.max_pool(_conv1, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME')
# padding = 'SAME' 代表会填补0, padding = 'VALID' 代表不会填补，丢弃最后的行列元素
_pool_dr1 = tf.nn.dropout(_pool1, _keepratio)
# conv layer2
_conv2 = tf.nn.conv2d(_pool_dr1, _w['wc2'], strides=[1,1,1,1], padding='SAME')
_conv2 = tf.nn.relu(tf.nn.bias_add(_conv2, _b['bc2']))
_pool2 = tf.nn.max_pool(_conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
# ksize = [batch_size_stride, height_stride, width_stride, channel_stride]
_pool_dr2 = tf.nn.dropout(_pool2, _keepratio)
# vectorize
_dense1 = tf.reshape(_pool_dr2, [-1, _w['wd1'].get_shape().as_list()[0]]) # reshape
# full connection layer
_fc1 = tf.nn.relu(tf.add(tf.matmul(_dense1, _w['wd1']), _b['bd1']))
_fc_dr1 = tf.nn.dropout(_fc1, _keepratio)
_out = tf.add(tf.matmul(_fc_dr1, _w['wd2']), _b['bd2'])
# result
out = {
'input_r': _input_r, 'conv1': _conv1, 'pool1': _pool1, 'pool_dr1': _pool_dr1,
'conv2': _conv2, 'pool2': _pool2, 'pool_dr2': _pool_dr2, 'dense1': _dense1,
'fc1': _fc1, 'fc_dr1': _fc_dr1, 'out': _out
}
return out
print('CNN ready')
# 网络定义
_pred = conv_basic(x, Weights, biases, keepratio)['out']
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=_pred, labels=y))
optm = tf.train.AdamOptimizer(learning_rate=0.001).minimize(cost)
_corr = tf.equal(tf.arg_max(_pred, 1), tf.arg_max(y, 1))
accr = tf.reduce_mean(tf.cast(_corr, tf.float32))
init = tf.global_variables_initializer()
# 保存模型
save_step = 1
saver = tf.train.Saver(max_to_keep=3)
print('graph ready')
do_train = 1
# 迭代计算
training_epochs = 30
batch_size = 20
display_step = 5
sess = tf.Session()
sess.run(init)
if do_train == 1:
for epoch in range(training_epochs):
avg_cost = 0
# total_batch = int(mnist.train.num_examples/batch_size)
total_batch = 5
for i in range(total_batch):
batch_xs, batch_ys = mnist.train.next_batch(batch_size)
sess.run(optm, feed_dict={x: batch_xs, y: batch_ys, keepratio: 0.5})
avg_cost += sess.run(cost, feed_dict={x: batch_xs, y: batch_ys, keepratio: 1.0})/total_batch
# 显示结果
if (epoch+1) % display_step == 0:
print('Epoch: %03d/%03d cost: %.9f' % (epoch+1, training_epochs, avg_cost))
feeds = {x: batch_xs, y: batch_ys, keepratio: 1.0}
train_acc = sess.run(accr, feed_dict=feeds)
print('Train accuracy: %.3f' % train_acc)
# 保存结果
if epoch % save_step == 0:
saver.save(sess, './data/nets/cnn_mnist_basic.ckpt-'+str(epoch+1))
do_train = 0
print('optmization finished')
if do_train == 0:
with tf.Session() as sess:
saver.restore(sess, './data/nets/cnn_mnist_basic.ckpt-' + str(training_epochs))
feeds = {x: mnist.test.images, y: mnist.test.labels, keepratio: 1.0}
test_acc = sess.run(accr, feed_dict=feeds)
print('Test accuracy: %.3f' % test_acc)
print('test finished')
通过简单对比，可以察觉，我们中相应地方添加了以下代码片段：
# 保存模型
save_step = 1
saver = tf.train.Saver(max_to_keep=3)
do_train = 1
并且将迭代部分的代码做如下修改：
if do_train == 1:
for epoch in range(training_epochs):
avg_cost = 0
# total_batch = int(mnist.train.num_examples/batch_size)
total_batch = 5
for i in range(total_batch):
batch_xs, batch_ys = mnist.train.next_batch(batch_size)
sess.run(optm, feed_dict={x: batch_xs, y: batch_ys, keepratio: 0.5})
avg_cost += sess.run(cost, feed_dict={x: batch_xs, y: batch_ys, keepratio: 1.0})/total_batch
# 显示结果
if (epoch+1) % display_step == 0:
print('Epoch: %03d/%03d cost: %.9f' % (epoch+1, training_epochs, avg_cost))
feeds = {x: batch_xs, y: batch_ys, keepratio: 1.0}
train_acc = sess.run(accr, feed_dict=feeds)
print('Train accuracy: %.3f' % train_acc)
# 保存结果
if epoch % save_step == 0:
saver.save(sess, './data/nets/cnn_mnist_basic.ckpt-'+str(epoch+1))
do_train = 0
print('optmization finished')
if do_train == 0:
with tf.Session() as sess:
saver.restore(sess, './data/nets/cnn_mnist_basic.ckpt-' + str(training_epochs))
feeds = {x: mnist.test.images, y: mnist.test.labels, keepratio: 1.0}
test_acc = sess.run(accr, feed_dict=feeds)
print('Test accuracy: %.3f' % test_acc)
print('test finished')
以上代码中主要添加了判断模型的训练过程和测试过程的if语句，并且通过模型的写入和读取方式重新加载绘话(Session)中参数值进行数据测试。
运行结果
通过30次的迭代，模型训练和测试结果如下：
Epoch: 005/030 cost: 2.127988672
Train accuracy: 0.450
Epoch: 010/030 cost: 1.514546847
Train accuracy: 0.600
Epoch: 015/030 cost: 1.691785026
Train accuracy: 0.600
Epoch: 020/030 cost: 1.580975151
Train accuracy: 0.750
Epoch: 025/030 cost: 1.391473675
Train accuracy: 0.800
Epoch: 030/030 cost: 1.286683488
Train accuracy: 0.750
optmization finished
Test accuracy: 0.744
test finished