TensorFlow入门(二) - mnist手写数字识别(模型保存加载)

# model.py
class Network:
    def __init__(self):
        # 记录已经训练的次数
        self.global_step = tf.Variable(0, trainable=False)

        # ... 中间省略网络结构

        # minimize 可传入参数 global_step， 每次训练 global_step的值会增加1
        # 因此，可以通过计算self.global_step这个张量的值，知道当前训练了多少步
        self.train = tf.train.GradientDescentOptimizer(0.001).minimize(
            self.loss, global_step=self.global_step)

CKPT_DIR = 'ckpt' # 定义模型存储的位置
net = Network()
sess = tf.Session()
sess.run(tf.global_variables_initializer())

# tf.train.Saver是用来保存训练结果的。
# max_to_keep 用来设置最多保存多少个模型，默认是5
# 如果保存的模型超过这个值，最旧的模型将被删除
saver = tf.train.Saver(max_to_keep=10)

train_step = 10000 # 总的训练次数10000
step = 0 # 记录训练次数, 初始化为0
save_interval = 1000 # 每隔1000步保存模型

while step < train_step:
    # ...省略训练代码

    step = sess.run(net.global_step)
    # 模型保存在ckpt文件夹下
    # 模型文件名最后会增加global_step的值，比如1000的模型文件名为 model-1000
    if step % save_interval == 0:
        saver.save(sess, CKPT_DIR + '/model', global_step=step)

CKPT_DIR = 'ckpt'
net = Network()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver(max_to_keep=10)

train_step = 10000
step = 0
save_interval = 1000

# 开始训练前，检查ckpt文件夹，看是否有checkpoint文件存在。
# 如果存在，则读取checkpoint文件指向的模型，restore到sess中。
ckpt = tf.train.get_checkpoint_state(CKPT_DIR)
if ckpt and ckpt.model_checkpoint_path:
    saver.restore(sess, ckpt.model_checkpoint_path)
    # 读取网络中的global_step的值，即当前已经训练的次数
    step = sess.run(net.global_step)
    print('Continue from')
    print('        -> Minibatch update : ', step)

while step < train_step:
    # ...省略训练代码

import numpy as np
from PIL import Image


class Predict:
    def __init__(self):
        self.net = Network()
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())
        self.restore() # 加载模型到sess中

    def restore(self):
        saver = tf.train.Saver()
        ckpt = tf.train.get_checkpoint_state(CKPT_DIR)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(self.sess, ckpt.model_checkpoint_path)
        else:
            raise FileNotFoundError("未保存任何模型")

    def predict(self, image_path):
        # ...省略

class Predict:
    # ...
    
    def predict(self, image_path):
        # 读图片并转为黑白的
        img = Image.open(image_path).convert('L')
        flatten_img = np.reshape(img, 784)
        x = np.array([1 - flatten_img])
        y = self.sess.run(self.net.y, feed_dict={self.net.x: x})

        # 因为x只传入了一张图片，取y[0]即可
        # np.argmax()取得独热编码最大值的下标，即代表的数字
        print(image_path)
        print('        -> Predict digit', np.argmax(y[0]))