|
|
import os
|
|
|
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()
|
|
|
# 图片归一化处理
|
|
|
train_images_scaled = train_images / 255
|
|
|
|
|
|
# 构建网络 开始
|
|
|
'''
|
|
|
卷积 池化 卷积 池化 全连接 全连接
|
|
|
'''
|
|
|
model = keras.Sequential()
|
|
|
model.add(keras.layers.Conv2D(8, (3, 3), activation='relu', input_shape=(28, 28, 1)))
|
|
|
model.add(keras.layers.MaxPooling2D(2, 2))
|
|
|
model.add(keras.layers.Conv2D(8, (3, 3), activation='relu'))
|
|
|
model.add(keras.layers.MaxPooling2D(2, 2))
|
|
|
|
|
|
model.add(keras.layers.Flatten()) # 扁平化处理，实现从卷积到全连接的过渡
|
|
|
model.add(keras.layers.Dense(128, activation=tf.nn.relu))
|
|
|
model.add(keras.layers.Dense(10, activation=tf.nn.softmax))
|
|
|
|
|
|
model.compile(optimizer='adam', loss=tf.losses.sparse_categorical_crossentropy, metrics=['accuracy'])
|
|
|
# 构建网络 结束
|
|
|
|
|
|
# 打印网络结构
|
|
|
model.summary()
|
|
|
|
|
|
# 回调函数，用于训练中保存模型
|
|
|
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
|
|
|
save_weights_only=True,
|
|
|
verbose=1)
|
|
|
|
|
|
# 发现检查点中存在保存的模型权值，则认为上次训练被中断，读取权值继续训练
|
|
|
if (os.path.exists(checkpoint_dir)):
|
|
|
print('检测到未完成的训练，已读取检查点权值继续训练')
|
|
|
model.load_weights(checkpoint_path)
|
|
|
|
|
|
# 训练
|
|
|
history = model.fit(
|
|
|
train_images_scaled.reshape(-1, 28, 28, 1),
|
|
|
train_labels,
|
|
|
epochs=8,
|
|
|
validation_data=(test_images.reshape(-1, 28, 28, 1), test_labels),
|
|
|
callbacks=[cp_callback]
|
|
|
)
|
|
|
|
|
|
# 保存模型
|
|
|
model.save(model_path)
|
|
|
|
|
|
# 删除检查点文件
|
|
|
shutil.rmtree(checkpoint_dir)
|
|
|
print("已删除检查点文件")
|
|
|
|
|
|
# 评估
|
|
|
results = model.evaluate(test_images.reshape(-1, 28, 28, 1), test_labels)
|
