首先介绍一下我的项目目录:
1.创建train.py,编写训练程序
import tensorflow as tf
import random
import matplotlib.pyplot as plt
def modelTrain():
# 导入数据集
mnist = tf.keras.datasets.mnist
(x_train_all, y_train_all), (x_test, y_test) = mnist.load_data()
# 数据预处理:将像素值缩放到 0 到 1 之间
x_train_all, x_test = x_train_all / 255.0, x_test / 255.0
# 随机选择20000张图片进行训练
train_indices = random.sample(range(0, len(x_train_all)), 20000)
x_train, y_train = x_train_all[train_indices], y_train_all[train_indices]
# 定义模型
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(10, activation='softmax')
])
# 编译模型
model.compile(optimizer='sgd',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
# 训练模型
history = model.fit(x_train, y_train, epochs=50,
validation_data=(x_test, y_test))
# 评估模型
test_loss, test_acc = model.evaluate(x_test, y_test)
print('Test accuracy:', test_acc)
# 绘制训练和验证时的准确率和损失
train_loss = history.history['loss']
val_loss = history.history['val_loss']
train_acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
# 绘制损失曲线
plt.subplot(2, 1, 1) # 设置子图,参数分别表示行数、列数、子图编号
plt.plot(train_loss, label='Training Loss')
plt.plot(val_loss, label='Validation Loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
# 绘制准确率曲线
plt.subplot(2, 1, 2) # 设置子图,参数分别表示行数、列数、子图编号
plt.plot(train_acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()
model.save('my_model.h5')
2.创建test.py,编写测试程序,传参为模型路径+测试图片路径,返回值为预测结果
import tensorflow as tf
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
def predict_custom_images(model_path, image_path):
# 加载模型
model = tf.keras.models.load_model(model_path)
# 打开图像并调整大小
img = Image.open(image_path)
img_resized = img.resize((28, 28))
# 将图像转换为灰度图像(如果不是的话)
if img_resized.mode != 'L':
img_resized = img_resized.convert('L')
# 输出灰度图像(调试语句)
plt.imshow(img_resized, cmap='gray')
plt.show()
# 转换为 numpy 数组
img_array = np.array(img_resized)
# 缩放像素值到 0-1 范围
img_array = img_array / 255.0
# 调整形状以适应模型
img_array = np.reshape(img_array, (1, 28, 28))
# 进行预测
pred = model.predict(img_array)
pred_label = np.argmax(pred)
return pred_label
3.编写主函数main.py,执行程序
from train import *
from test import predict_custom_images
# 训练模型并保存模型
modelTrain()
# 测试自定义手写数字图片
model_path = 'my_model.h5'
for i in range(10):
img_path = 'img{}.png'.format(i)
pred_label = predict_custom_images(model_path, img_path)
print('{} 的预测结果是:{}'.format(img_path, pred_label))
运行结果:
1.训练中:
2.图表:
(这里结束之后,会多一个模型文件:my_model.h5,用于接下来的测试)
3.测试结果:
在这里大家一定发现了,其中“7”和“9”的预测结果错误,毕竟7和1,9和7长得还是比较像的
我想大概是:
因为每次只训练了20000个样本,数量比较小,而且只是50轮次,就出现了这种情况
大家可以通过修改代码中的轮次和每次训练数(不超过60000)