1
2
3
4
5
6
7
8
9
10
11

|--mnist/
  |--data_set/ 训练以及测试数据集
  |--test_images/  多张测试图片
    |--0.png
    |--1.png
    |--4.png
  |--v2/  
    |--ckpt/ 模型保存在这里！！！
    |--model.py  网络模型
    |--train.py  训练代码
    |--predict.py  预测代码

1
2
3
4
5
6
7
8
9
10
11
12

class Network:
    def __init__(self):
        
        self.global_step = tf.Variable(0, trainable=False)

        

        
        
        self.train = tf.train.GradientDescentOptimizer(0.001).minimize(
            self.loss, global_step=self.global_step)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

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 

while step < train_step:
    

    step = sess.run(net.global_step)
    
    
    if step % save_interval == 0:
        saver.save(sess, CKPT_DIR + '/model', global_step=step)

1
2
3

model_checkpoint_path: "model-2000"
all_model_checkpoint_paths: "model-1000"
all_model_checkpoint_paths: "model-2000"

1
2
3
4
5
6
7
8
9
10
11
12

|--v2/  
  |--ckpt/ 模型保存在这里！！！
    |--checkpoint
    |--model-1000.data-00000-of-00001
    |--model-1000.index
    |--model-1000.meta
    |--model-2000.data-00000-of-00001
    |--model-2000.index
    |--model-2000.meta
  |--model.py  网络模型
  |--train.py  训练代码
  |--predict.py  预测代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

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 = tf.train.get_checkpoint_state(CKPT_DIR)
if ckpt and ckpt.model_checkpoint_path:
    saver.restore(sess, ckpt.model_checkpoint_path)
    
    step = sess.run(net.global_step)
    print('Continue from')
    print('        -> Minibatch update : ', step)

while step < train_step:
    

1
2
3

Continue from
        -> Minibatch update :  2000
第 3000步，...

1
2
3

Continue from
        -> Minibatch update :  1000
第 2000步，...

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

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() 

    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):
        

1
2
3
4
5
6
7
8
9
10
11
12
13
14

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})

        
        
        print(image_path)
        print('        -> Predict digit', np.argmax(y[0]))