|
|
import os
|
|
|
import tensorflow as tf
|
|
|
import matplotlib.pyplot as plt
|
|
|
from tensorflow import keras
|
|
|
import numpy as np
|
|
|
|
|
|
# 抑制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)
|
|
|
|
|
|
# 构建网络 开始
|
|
|
'''
|
|
|
卷积 池化 卷积 池化 全连接 全连接
|
|
|
'''
|
|
|
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()
|
|
|
|
|
|
# 回调函数，用于训练中保存模型
|
|
|
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,
|
|
|
verbose=1)
|
|
|
|
|
|
# 图片归一化处理
|
|
|
train_images_scaled = train_images / 255
|
|
|
|
|
|
# 训练
|
|
|
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)
|
|
|
|
|
|
# 评估
|
|
|
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()
|
