TensorFlow可以保存训练过的模型，不仅在训练过程中断后，可以继续上次训练过程；还可以进行迁移学习，在别人的训练的模型基础上训练自己的模型。可谓十分方便。
TensorFlow保存模型checkpoint后生成以下文件:
|—checkpoint
|—model_name.data-00000-of-00001
|—model_name.index
|—model_name.meta
model_name为定义好的模型名字
model_name.meta为图文件
model_name.data为数据文件
- 保存模型
saver = tf.train.Saver() #创建saver对象
saver.save(sess, checkpoint_path) #将sess保存到定义好的路径下- 恢复模型
只加载数据
saver.restore(sess, checkpoint_path) #从路径中恢复模型到会话sess加载图和数据
meta_path = 'model_name.meta' #图路径
model_path = 'model_name' #模型路径
saver = tf.train.import_meta_graph(meta_path) #加载图
with tf.Session() as sess:
saver.restore(sess, model_path) #恢复会话sess并加载数据
graph = tf.get_default_graph()
x = graph.get_tensor_by_name('InputData:0') #从图中获取tensor所有程序代码(基于TensorFlow基础教程:搭建简单的DNN实现手写数字识别)改写
训练代码
# coding: utf-8
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
import tensorflow as tf
learning_rate = 0.001
train_epochs = 10
batch_size = 64
checkpoint_path = 'checkpoint/'
n_input = 784
n_hidden1 = 100
n_hidden2 = 100
n_classes = 10
#name参数，记录变量名字
x = tf.placeholder(tf.float32, shape=[None, n_input], name='InputData')
y = tf.placeholder(tf.float32, shape=[None, n_classes], name='LabelData')
weights = {'w1': tf.Variable(tf.random_normal([n_input, n_hidden1]), name='W1'),
'w2': tf.Variable(tf.random_normal([n_hidden1, n_hidden2]), name='W2'),
'w3': tf.Variable(tf.random_normal([n_hidden2, n_classes]), name='W3')}
biases = {'b1': tf.Variable(tf.random_normal([n_hidden1]), name='b1'),
'b2': tf.Variable(tf.random_normal([n_hidden2]), name='b2'),
'b3': tf.Variable(tf.random_normal([n_classes]), name='b3')}
def inference(input_x):
layer_1 = tf.nn.relu(tf.matmul(x, weights['w1']) + biases['b1'])
layer_2 = tf.nn.relu(tf.matmul(layer_1, weights['w2']) + biases['b2'])
out_layer = tf.matmul(layer_2, weights['w3']) + biases['b3']
return out_layer
#定义计算过程的名字
with tf.name_scope('Inference'):
logits = inference(x)
with tf.name_scope('Loss'):
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=y))
with tf.name_scope('Optimizer'):
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss)
with tf.name_scope('Accuracy'):
pre_correct = tf.equal(tf.argmax(y, 1), tf.argmax(tf.nn.softmax(logits), 1))
accuracy = tf.reduce_mean(tf.cast(pre_correct, tf.float32), name='acc')
print(accuracy)
init = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
total_batch = int(mnist.train.num_examples / batch_size)
checkpoint = tf.train.get_checkpoint_state(checkpoint_path) #获取checkpoint状态
if checkpoint and checkpoint.model_checkpoint_path:
saver.restore(sess, checkpoint_path+'model.ckpt') #加载数据
print('continue last train!!')
else:
print('restart train!!')
for epoch in range(train_epochs):
for batch in range(total_batch):
batch_x, batch_y = mnist.train.next_batch(batch_size)
sess.run(train_op, feed_dict={x:batch_x, y:batch_y})
if (epoch+1) % 5 == 0:
loss_, acc = sess.run([loss, accuracy], feed_dict={x:batch_x, y:batch_y})
print("epoch {}, loss {:.4f}, acc {:.3f}".format(epoch, loss_, acc))
saver.save(sess, checkpoint_path+'model.ckpt') #模型名字model.ckpt
print("optimizer finished!")
print("模型保存在"， checkpoint_path)模型恢复并计算测试集准确度
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
import tensorflow as tf
meta_path = 'checkpoint/model.ckpt.meta' #图路径
model_path = 'checkpoint/model.ckpt' #数据路径
saver = tf.train.import_meta_graph(meta_path) #加载图
with tf.Session() as sess:
saver.restore(sess, model_path) #加载数据
graph = tf.get_default_graph()
x = graph.get_tensor_by_name('InputData:0') #加载张量
y = graph.get_tensor_by_name('LabelData:0')
accuracy = graph.get_tensor_by_name('Accuracy/acc:0')
#计算测试集的准确度
test_acc = sess.run(accuracy, feed_dict={x:mnist.test.images, y:mnist.test.labels})
print('test accuracy', test_acc)github源码下载
https://github.com/gamersover/tensorflow_basic_tutorial/tree/master/model_save_tutorial