TensorFlow进行MNIST数据集手写数字识别，保存模型并且进行外部手写图片测试

news2024/12/25 15:30:32

首先，你已经配置好Anaconda3的环境，下载了TensorFlow模块，并且会使用jupyter了，那么接下来就是MNIST实验步骤。

数据集官网下载：MNIST handwritten digit database, Yann LeCun, Corinna Cortes and Chris Burgeshttp://yann.lecun.com/exdb/mnist/

将上面四个全部下载，都是数据集，其中前两个是训练集，后两个是测试集

当然上面的数据集使用TensorFlow1.x要方便一点，本文章是使用TensorFlow2.x版本，使用npz格式数据集

数据集官网下载：https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz

导入该导的包

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import datetime

接下来是读取数据集，有以下方法可以读取

第一种非常简单，直接使用tensorflow中的方法从官网下载数据集，这样就不用手动去官网下载数据集了，但是比较耗时间

mnist=tf.keras.datasets.mnist

#获取数据，训练集，测试集 60k训练，10K测试
#网络下载
(x_train,y_train),(x_test,y_test)=mnist.load_data()

第二种方法，使用下载好的本地npz文件，将下面的路径改成你自己的

#本地加载
dataset = np.load('/home/tensor/jupyter/MNIST_data2/mnist.npz')
# 获取训练集和测试集
x_train, y_train = dataset['x_train'], dataset['y_train']
x_test, y_test = dataset['x_test'], dataset['y_test']

下一步，原图片的大小为28x28，将图片转成32x32的大小效果更好

#首先是数据 INPUT 层，输入图像的尺寸统一归一化为32*32。
x_train= np.pad(x_train,((0,0),(2,2),(2,2)),'constant',constant_values=0) #28*28-》32*32
x_test= np.pad(x_test,((0,0),(2,2),(2,2)),'constant',constant_values=0) #28*28-》32*32
#print(x_train.shape,x_test.shape)

#数据集格式转换
# x_train=x_train.astype('float32')
# x_train=x_train.astype('float32')

#归一化，就是为了限定你的输入向量的最大值跟最小值不超过你的隐层跟输出层函数的限定范围。
x_train=x_train/255#归一化
x_test=x_test/255#归一化

x_train=x_train.reshape(x_train.shape[0],32,32,1)
x_test=x_test.reshape(x_test.shape[0],32,32,1)
print(x_train.shape,x_test.shape)

输出结果如下

创建神经网络模型，这里我简单构建一个卷积神经网络模型

def LeNetModel():
#模型实例化,根据LeNet 的结构
    model=tf.keras.models.Sequential([
        tf.keras.layers.Conv2D(filters=6,kernel_size=(5,5),padding='valid',activation=tf.nn.relu,input_shape=(32,32,1)),
        tf.keras.layers.AveragePooling2D(pool_size=(2,2),strides=(2,2),padding='same'),
        tf.keras.layers.Conv2D(filters=16,kernel_size=(5,5),padding='valid',activation=tf.nn.relu,input_shape=(32,32,1)),
        tf.keras.layers.AveragePooling2D(pool_size=(2,2),strides=(2,2),padding='same'),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(units=120,activation=tf.nn.relu),
        tf.keras.layers.Dense(units=84,activation=tf.nn.relu),
        tf.keras.layers.Dense(units=10,activation=tf.nn.softmax),
        ])
    return model

加载模型，并且输出摘要

model = LeNetModel()
model.summary() #输出摘要

输出如下：

定义超参数

num_epochs=1#训练次数
batch_size=60#每个批次喂多少张图片
lr=0.001#学习率

开始训练模型

#优化器
adam_optimizer=tf.keras.optimizers.Adam(lr)
     
model.compile(
        optimizer=adam_optimizer,
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=['accuracy']
    )

start_time=datetime.datetime.now() #开始训练时间
     
model.fit(x=x_train,y=y_train,batch_size=batch_size,epochs=num_epochs)
end_time=datetime.datetime.now() #训练结束时间
time_cost=end_time-start_time #训练总时间
print('time_cost: ',time_cost)
model.save('leNet_model.h5') #保存模型
print(model.evaluate(x_test,y_test))
print("Finished!")

输出结果如下：

接下来选择一张测试集图片或多张测试集图片进行测试

def pred_function(image,label):
    image_index=20
     
    # 预测
    pred=model.predict(image.reshape(1,32,32,1))
    print("label:",label,"predict result:",pred.argmax())
     
    # 显示
    plt.imshow(image.reshape(32,32))
    plt.savefig("predict_num.jpg")
    plt.show()

index = 100 #选择第100张图片进行测试
pred_function(x_test[index],y_test[index])

结果如下：

在结果中，label表示正确值，predict result表示模型测试的结果值

完成以上步骤，那么下面就是使用已保存的模型来预测我们自己手写的数字，将10.jpg改成你的手写图片路径

#模型使用
import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib as m
import numpy as np
import cv2
import os

mnist=tf.keras.datasets.mnist
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# np.set_printoptions(threshold=np.inf)
#加载模型
def digit_predict():
    model=tf.keras.models.load_model('leNet_model.h5')
     
    #图片预处理
    img=cv2.imread('10.jpg')
    print(img.shape)
    plt.imshow(img)
    plt.show()
     
    #灰度图
    img=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    print(img.shape)
    plt.imshow(img,cmap='Greys')
    plt.show()
     
    #取反
    img=cv2.bitwise_not(img)
    plt.imshow(img,cmap='Greys')
    plt.show()
     
    print('二值化前：',img.shape)
    print('二值化前：',img)
    #纯黑 纯白 二值化
    img[img<=100]=0
    img[img>=140]=255
    plt.imshow(img,cmap='Greys')
    plt.show()
    print('二值化后：',img.shape)
    print('二值化后：',img)
     
    #尺寸
    img=cv2.resize(img,(32,32))
    print('尺寸:',img.shape)
    print('尺寸',img)
     
    #归一化
    img=img/255
    print('归一化:',img.shape)
    print('归一化:',img)#0和1组成

    #预测
    pred=model.predict(img.reshape(1,32,32,1))
    print('prediction Number: ',pred.argmax())
     
    #打印图片信息
    plt.imshow(img)
    plt.show()
     
digit_predict()

下面是测试我自己手写数字的结果：