手写数字识别——KNN模型实现

MNIST手写数字识别

        MNIST 手写数字数据库有一个包含 60,000 个示例的训练集和一个包含 10,000 个示例的测试集。

• 每个图像高 28 像素，宽28 像素，共784个像素。

• 每个像素取值范围[0,255]，取值越大意味着该像素颜色越深

        下载：http://yann.lecun.com/exdb/mnist/

import os
from torchvision import datasets

# 设置数据集的根目录
os.environ['TORCH_HOME'] = './MNIST'

# 下载数据集
train_dataset = datasets.MNIST(root=os.getenv('TORCH_HOME'), train=True, download=True,)
test_dataset = datasets.MNIST(root=os.getenv('TORCH_HOME'), train=False, download=True,)

print(train_dataset.train_data.shape)
print(train_dataset.train_labels.shape)
print(test_dataset.test_data.shape)
print(test_dataset.test_labels.shape)

torch.Size([60000, 28, 28])
torch.Size([60000])
torch.Size([10000, 28, 28])
torch.Size([10000])

KNN算法预测

import matplotlib.pyplot as plt
from tensorflow.keras.datasets import mnist
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import joblib
import numpy as np

# 加载MNIST数据集
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 数据归一化
train_images = train_images / 255.0
test_images = test_images /255.0

# 将数据全部组合
x = np.concatenate((train_images,test_images),axis=0).reshape(70000, -1)
y = np.concatenate((train_labels, test_labels),axis=0)

# 数据集划分
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, stratify=y, random_state=0)

# 定义并训练KNN分类器
knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(x_train, y_train)

print('测试集准确率：',knn.score(x_test,y_test))

# 保存训练好的KNN模型
joblib.dump(knn, './knn_model.pth')

# 加载并预测新图像
img = plt.imread('./test_image.png')
plt.imshow(img)

# 加载保存的KNN模型
knn_loaded = joblib.load('./knn_model.pth')

# 对图像进行预测
img_flat = img.reshape(1, -1)  # 将图像展平成一维数组
y_predict = knn_loaded.predict(img_flat)
print(f'Predicted Label: {y_predict}')

测试集准确率： 0.9735
Predicted Label: [2]