训练与预测分离，加入断点训练

predict.py

@@ -4,13 +4,13 @@ import matplotlib.pyplot as plt
 from tensorflow.keras.models import load_model
 import numpy as np
 
+# 模型路径
+model_path = "./myModel/myModel.h5"
+
 # 抑制tensorflow，以防显存占用过多报错
 config = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True))
 sess = tf.compat.v1.Session(config=config)
 
-checkpoint_path = "./training_2/cp.ckpt"
-model_path = "./myModel/myModel.h5"
-
 # 读取手写数字数据集
 num_mnist = tf.keras.datasets.mnist
 (train_images, train_labels), (test_images, test_labels) = num_mnist.load_data()
@@ -25,7 +25,8 @@ model.summary()
 model.evaluate(test_images.reshape(-1, 28, 28, 1), test_labels)
 
 # 可视化预测效果
-testShow = test_labels[:100]
+show_num = 300
+testShow = test_labels[:show_num]
 
 pred = model.predict(test_images.reshape(-1, 28, 28, 1))
 predict = []
@@ -35,7 +36,7 @@ for item in pred:
 plt.figure()
 plt.title('Conv Predict')
 plt.ylabel('number')
-plt.plot(range(testShow.size), predict[:100], label='predict')
+plt.plot(range(testShow.size), predict[:show_num], label='predict')
 plt.plot(range(testShow.size), testShow, label='result')
 plt.legend()
 plt.show()

train.py

@@ -3,22 +3,27 @@ import tensorflow as tf
 import matplotlib.pyplot as plt
 from tensorflow import keras
 import numpy as np
+import shutil
+
+# 检查点路径与模型保存路径
+checkpoint_path = "./checkpoint/cp.ckpt"
+model_path = "./myModel/myModel.h5"
+checkpoint_dir = os.path.dirname(checkpoint_path)
+model_dir = os.path.dirname(model_path)
 
 # 抑制tensorflow，以防显存占用过多报错
 config = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True))
 sess = tf.compat.v1.Session(config=config)
+
 # 读取手写数字数据集
 num_mnist = keras.datasets.mnist
 (train_images, train_labels), (test_images, test_labels) = num_mnist.load_data()
-# 打印shape
-print(train_images.shape)
-print(train_labels.shape)
-print(test_images.shape)
-print(test_labels.shape)
+# 图片归一化处理
+train_images_scaled = train_images / 255
 
 # 构建网络 开始
 '''
-卷积 池化 卷积 池化  全连接 全连接
+卷积 池化 卷积 池化 全连接 全连接
 '''
 model = keras.Sequential()
 model.add(keras.layers.Conv2D(8, (3, 3), activation='relu', input_shape=(28, 28, 1)))
@@ -37,16 +42,14 @@ model.compile(optimizer='adam', loss=tf.losses.sparse_categorical_crossentropy,
 model.summary()
 
 # 回调函数，用于训练中保存模型
-checkpoint_path = "./training_2/cp.ckpt"
-model_path = "./myModel/myModel.h5"
-checkpoint_dir = os.path.dirname(checkpoint_path)
-
 cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
-                                                 #save_weights_only=True,
+                                                 save_weights_only=True,
                                                  verbose=1)
 
-# 图片归一化处理
-train_images_scaled = train_images / 255
+# 发现检查点中存在保存的模型权值，则认为上次训练被中断，读取权值继续训练
+if (os.path.exists(checkpoint_dir)):
+    print('检测到未完成的训练，已读取检查点权值继续训练')
+    model.load_weights(checkpoint_path)
 
 # 训练
 history = model.fit(
@@ -57,24 +60,12 @@ history = model.fit(
     callbacks=[cp_callback]
 )
 
-#保存模型
+# 保存模型
 model.save(model_path)
 
+# 删除检查点文件
+shutil.rmtree(checkpoint_dir)
+print("已删除检查点文件")
+
 # 评估
 results = model.evaluate(test_images.reshape(-1, 28, 28, 1), test_labels)
-
-# 可视化预测效果
-testShow = test_labels[:100]
-
-pred = model.predict(test_images.reshape(-1, 28, 28, 1))
-predict = []
-for item in pred:
-    predict.append(np.argmax(item))
-
-plt.figure()
-plt.title('Conv Predict')
-plt.ylabel('number')
-plt.plot(range(testShow.size), predict[:100], label='predict')
-plt.plot(range(testShow.size), testShow, label='result')
-plt.legend()
-plt.show()

training_1/checkpoint

@@ -1,2 +0,0 @@
-model_checkpoint_path: "cp.ckpt"
-all_model_checkpoint_paths: "cp.ckpt"