TensorFlow学习4：保存和读取 TensorFlow 模型
训练一个模型的时间很长。但是你一旦关闭了 TensorFlow session，你所有训练的权重和偏置项都丢失了。如果你计划在之后重新使用这个模型，你需要重新训练！
幸运的是，TensorFlow 可以让你通过一个叫 tf.train.Saver 的类把你的进程保存下来。这个类可以把任何 tf.Variable存到你的文件系统。
保存变量
让我们通过一个简单地例子来保存 weights 和 bias Tensors。第一个例子你只是存两个变量，后面会教你如何把一个实际模型的所有权重保存下来。
import tensorflow as tf
# The file path to save the data
# 文件保存路径
save_file = './model.ckpt'
# Two Tensor Variables: weights and bias
# 两个 Tensor 变量：权重和偏置项
weights = tf.Variable(tf.truncated_normal([2, 3]))
bias = tf.Variable(tf.truncated_normal([3]))
# Class used to save and/or restore Tensor Variables
# 用来存取 Tensor 变量的类
saver = tf.train.Saver()
with tf.Session() as sess:
# Initialize all the Variables
# 初始化所有变量
sess.run(tf.global_variables_initializer())
# Show the values of weights and bias
# 显示变量和权重
print('Weights:')
print(sess.run(weights))
print('Bias:')
print(sess.run(bias))
# Save the model
# 保存模型
saver.save(sess, save_file)Weights:
[[-0.97990924 1.03016174 0.74119264]
[-0.82581609 -0.07361362 -0.86653847]]
Bias:
[ 1.62978125 -0.37812829 0.64723819]
weights 和 bias Tensors 用 tf.truncated_normal() 函数设定了随机值。用 tf.train.Saver.save() 函数把这些值被保存在save_file 位置，命名为 "model.ckpt"，（".ckpt" 扩展名表示"checkpoint"）。
如果你使用 TensorFlow 0.11.0RC1 或者更新版，一个叫做 "model.ckpt.meta" 的文件也会生成。它包含了 TensorFlow graph。
加载变量
现在这些变量已经存好了，让我们把它们加载到新模型里。
# Remove the previous weights and bias
# 移除之前的权重和偏置项
tf.reset_default_graph()
# Two Variables: weights and bias
# 两个变量：权重和偏置项
weights = tf.Variable(tf.truncated_normal([2, 3]))
bias = tf.Variable(tf.truncated_normal([3]))
# Class used to save and/or restore Tensor Variables
# 用来存取 Tensor 变量的类
saver = tf.train.Saver()
with tf.Session() as sess:
# Load the weights and bias
# 加载权重和偏置项
saver.restore(sess, save_file)
# Show the values of weights and bias
# 显示权重和偏置项
print('Weight:')
print(sess.run(weights))
print('Bias:')
print(sess.run(bias))Weights:
[[-0.97990924 1.03016174 0.74119264]
[-0.82581609 -0.07361362 -0.86653847]]
Bias:
[ 1.62978125 -0.37812829 0.64723819]
保存一个训练好的模型
让我们看看如何训练一个模型并保存它的权重。
从一个模型开始：
# Remove previous Tensors and Operations
# 移除之前的 Tensors 和运算
tf.reset_default_graph()
from tensorflow.examples.tutorials.mnist import input_data
import numpy as np
learning_rate = 0.001
n_input = 784 # MNIST 数据输入 (图片尺寸: 28*28)
n_classes = 10 # MNIST 总计类别 (数字 0-9)
# Import MNIST data
# 加载 MNIST 数据
mnist = input_data.read_data_sets('.', one_hot=True)
# Features and Labels
# 特征和标签
features = tf.placeholder(tf.float32, [None, n_input])
labels = tf.placeholder(tf.float32, [None, n_classes])
# Weights & bias
# 权重和偏置项
weights = tf.Variable(tf.random_normal([n_input, n_classes]))
bias = tf.Variable(tf.random_normal([n_classes]))
# Logits - xW + b
logits = tf.add(tf.matmul(features, weights), bias)
# Define loss and optimizer
# 定义损失函数和优化器
cost = tf.reduce_mean(\
tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels))
optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)\
.minimize(cost)
# Calculate accuracy
# 计算准确率
correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))import math
save_file = './train_model.ckpt'
batch_size = 128
n_epochs = 100
saver = tf.train.Saver()
# Launch the graph
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# Training cycle
for epoch in range(n_epochs):
total_batch = math.ceil(mnist.train.num_examples / batch_size)
# Loop over all batches
for i in range(total_batch):
batch_features, batch_labels = mnist.train.next_batch(batch_size)
sess.run(
optimizer,
feed_dict={features: batch_features, labels: batch_labels})
# Print status for every 10 epochs
if epoch % 10 == 0:
valid_accuracy = sess.run(
accuracy,
feed_dict={
features: mnist.validation.images,
labels: mnist.validation.labels})
print('Epoch {:<3} - Validation Accuracy: {}'.format(
epoch,
valid_accuracy))
# Save the model
saver.save(sess, save_file)
print('Trained Model Saved.')Epoch 0 - Validation Accuracy: 0.06859999895095825
Epoch 10 - Validation Accuracy: 0.20239999890327454
Epoch 20 - Validation Accuracy: 0.36980000138282776
Epoch 30 - Validation Accuracy: 0.48820000886917114
Epoch 40 - Validation Accuracy: 0.5601999759674072
Epoch 50 - Validation Accuracy: 0.6097999811172485
Epoch 60 - Validation Accuracy: 0.6425999999046326
Epoch 70 - Validation Accuracy: 0.6733999848365784
Epoch 80 - Validation Accuracy: 0.6916000247001648
Epoch 90 - Validation Accuracy: 0.7113999724388123
Trained Model Saved.
加载训练好的模型
让我们从磁盘中加载权重和偏置项，验证测试集准确率
saver = tf.train.Saver()
# Launch the graph
with tf.Session() as sess:
saver.restore(sess, save_file)
test_accuracy = sess.run(
accuracy,
feed_dict={features: mnist.test.images, labels: mnist.test.labels})
print('Test Accuracy: {}'.format(test_accuracy))