Pytorch手写数字实验（自己手写识别）

一、概述

该文是基于PyTorch框架，采用CNN卷积神经网络实现手写数字识别，共采用了2个卷积层、两个池化层和三个线性层。（仅使用GPU进行训练，用的ide是pycharm）

本人也是初学者，如果有不对的地方，欢迎各位大佬提出意见和改进。

导入相关的包：

import time
import torch
from torch.utils.data import DataLoader
from torch import nn
import torchvision

二、构建网络模型

首先需要了解cnn网络模型框架：

比如传入的数字是‘5’，它的维度为（1，28，28）

首先通过一个5x5的卷积核（卷积层），其通道从1变为10，尺寸从28*28变为24*24.

在通过一个2x2卷积核（最大池化层），其通道没有变，为10，尺寸24*24变为12*12

然后通过一个5x5的卷积核（卷积层），其通道从10变为20，尺寸12*12变为8*8

再通过一个2x2卷积核（最大池化层），其通道没有变，为10，尺寸8*8变为4*4

最后通过一个线性层，通道从320变为10

填充边界----padding 和 步长----stride 通过下面这个公式进行计算

构建网络的代码为：

class Mnist(nn.Module):
    def __init__(self):
        super(Mnist, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 10, 5, 1, 0)
            , nn.MaxPool2d(kernel_size=(2, 2), padding=0, stride=2)
            , nn.Conv2d(10, 20, 5, 1, 0)
            , nn.MaxPool2d(kernel_size=(2, 2), padding=0, stride=2)
            , nn.Flatten()  # 展平
            , nn.Linear(320, 128) # 320是展平后的结果，20*4*4
            , nn.Linear(128, 64)
            , nn.Linear(64, 10)
            , nn.ReLU()
        )
    
    def forward(self, x):
        x = self.model(x)
        return x

三、数据集的导入

通过torchvision导入Mnists数据集，代码如下：

def loader_data(test_data_train):
    data_set = torchvision.datasets.MNIST('../dataset', test_data_train
                                          , torchvision.transforms.ToTensor()
                                          , download=True)  # 加载数据以及张量化
    return DataLoader(data_set, batch_size=32, shuffle=True)  # 每批32个样本随机打乱

为了简化代码和方便，将导入数据集封装成一个函数。

四、训练

将模型实例化和GPU设定

为了节省训练的时间，采用GPU进行训练。

device = 'cuda'
mnist = Mnist()
mnist.to(device)

采用GPU训练前，首先需要确定是否下载了CUDA驱动,如果没有安装，提前安装好。安装教程请参考：https://blog.csdn.net/qq_35831906/article/details/134349866?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522172423115416800178555617%2522%252C%2522scm%2522%253A%252220140713.130102334..%2522%257D&request_id=172423115416800178555617&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~top_positive~default-1-134349866-null-null.142^v100^pc_search_result_base1&utm_term=cuda%20pytorch&spm=1018.2226.3001.4187

训练代码如下：

def train(mnist, epoch):  # epoch训练轮数
    best_acc = 0
    train_data = loader_data(test_data_train=True) # 载入训练集
    test_data = loader_data(test_data_train=False) # 载入测试集
    start_time = time.time() # 这里是为了方便查看训练用了多久时间
    print("initial_accuracyL:{}".format(estimate(test_data, mnist)[1]))
    optim = torch.optim.Adam(mnist.parameters(), lr=0.001) # 采用Adam优化器，学习率为1
    for i in range(epoch):
        print('-----第{}轮训练开始-----'.format(i + 1))
        for data in train_data:
            images, targets = data
            images = images.to(device)  
            targets = targets.to(device)  
            mnist.zero_grad()  # 梯度清零
            outPuts = mnist.forward(images.view(-1, 1, 28, 28))
            loss_ = nn.CrossEntropyLoss()  # 交叉熵损失
            loss_.to(device)
            loss = loss_(outPuts, targets)
            loss.backward()  # 反向传播
            optim.step()     # 参数优化
        end_time = time.time()
        total_time = end_time - start_time
        print("accuracy:{}, 用时：{}min{:.2f}s, Loss{}".format(estimate(test_data, mnist)[1]
                                                        , total_time // 60
                                                        , total_time % 60
                                                        , estimate(test_data, mnist)[0]
                                                        ))
        epoch_acc = estimate(test_data, mnist)[1]  
        if epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model = mnist.state_dict()  # 这里是为了选出最好的模型
    
    torch.save(best_model, "Mnist.pth") # 保存模型

Adam（Adaptive Moment Estimation）是一种常用的优化算法，特别适用于训练神经网络和深度学习模型。它是一种自适应学习率的优化算法，可以根据不同参数的梯度信息来动态调整学习率，以提高训练的效率和稳定性。

训练整体代码

"""
author: XiaoShu
date: 2024-08-21
"""

import time
import torch
from torch.utils.data import DataLoader
from torch import nn
import torchvision

device = 'cuda'


class Mnist(nn.Module):
    def __init__(self):
        super(Mnist, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 10, 5, 1, 0)
            , nn.MaxPool2d(kernel_size=(2, 2), padding=0, stride=2)
            , nn.Conv2d(10, 20, 5, 1, 0)
            , nn.MaxPool2d(kernel_size=(2, 2), padding=0, stride=2)
            , nn.Flatten()
            , nn.Linear(320, 128)
            , nn.Linear(128, 64)
            , nn.Linear(64, 10)
            , nn.ReLU()
        )
    
    def forward(self, x):
        x = self.model(x)
        return x


mnist = Mnist()
mnist.to(device)


def loader_data(test_data_train):
    data_set = torchvision.datasets.MNIST('../dataset', test_data_train
                                          , torchvision.transforms.ToTensor(), download=True)  # 加载数据以及张量化
    return DataLoader(data_set, batch_size=32, shuffle=True)  # 每批32个样本随机打乱


# 预测模型性能
def estimate(test_data, mnist):
    yes_num = 0
    total_num = 0
    with torch.no_grad():
        for data in test_data:
            images, targets = data
            images = images.to(device)
            targets = targets.to(device)
            outPuts = mnist(images.view(-1, 1, 28, 28))
            loss_ = nn.CrossEntropyLoss()
            loss_.to(device)
            loss = loss_(outPuts, targets)
            outPuts = mnist.forward(images.view(-1, 1, 28, 28))
            for i, outPut in enumerate(outPuts):
                if torch.argmax(outPut) == targets[i]:
                    yes_num += 1
                total_num += 1
        accuracy = yes_num / total_num
    return loss, accuracy


def train(mnist, epoch):
    best_acc = 0
    train_data = loader_data(test_data_train=True)
    test_data = loader_data(test_data_train=False)
    start_time = time.time()
    print("initial_accuracyL:{}".format(estimate(test_data, mnist)[1]))
    optim = torch.optim.Adam(mnist.parameters(), lr=0.001)
    for i in range(epoch):
        print('-----第{}轮训练开始-----'.format(i + 1))
        for data in train_data:
            images, targets = data
            images = images.to(device)
            targets = targets.to(device)
            mnist.zero_grad()
            outPuts = mnist.forward(images.view(-1, 1, 28, 28))
            loss_ = nn.CrossEntropyLoss()
            loss_.to(device)
            loss = loss_(outPuts, targets)
            loss.backward()
            optim.step()
        end_time = time.time()
        total_time = end_time - start_time
        print("accuracy:{}, 用时：{}min{:.2f}s, Loss{}".format(estimate(test_data, mnist)[1]
                                                        , total_time // 60
                                                        , total_time % 60
                                                        , estimate(test_data, mnist)[0]
                                                        ))
        epoch_acc = estimate(test_data, mnist)[1]
        if epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model = mnist.state_dict()
    
    torch.save(best_model, "Mnist1.pth")


if __name__ == '__main__':
    train(mnist, 100)

五、模型性能评估

代码：

def estimate(test_data, mnist):
    yes_num = 0
    total_num = 0
    with torch.no_grad():
        for data in test_data:
            images, targets = data
            images = images.to(device)
            targets = targets.to(device)
            outPuts = mnist(images.view(-1, 1, 28, 28))
            loss_ = nn.CrossEntropyLoss()
            loss_.to(device)
            loss = loss_(outPuts, targets)
            outPuts = mnist.forward(images.view(-1, 1, 28, 28))
            for i, outPut in enumerate(outPuts):
                if torch.argmax(outPut) == targets[i]:
                    yes_num += 1
                total_num += 1
        accuracy = yes_num / total_num
    return loss, accuracy

六、训练后效果

本文训练了100轮

模型训练好后可以开始测试了，新建一个测试的py文件

测试代码(可视化)：

"""
author: XiaoShu
date: 2024-08-21
"""
import torchvision

import train
import torch
from PIL import Image
import matplotlib.pyplot as plt

img_path = '../test_two.png'
image = Image.open(img_path)
image1 = Image.open(img_path)
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((28, 28))
                                               , torchvision.transforms.ToTensor()])

image = transform(image)
image = torch.reshape(image, (-1, 1, 28, 28))
model = train.Mnist()
model.load_state_dict(torch.load('Mnist1.pth', map_location=torch.device('cpu')))
model.eval()
with torch.no_grad():
    outPut = model(image)
    prediction = outPut.argmax().item()
print('识别结果：{}'.format(prediction))

# 可视化
f, a = plt.subplots(1, 1, figsize=(7, 7))
a.set_title('predict:{}'.format(prediction))
a.axis('off')
a.imshow(image1)
plt.show()

我导入图片是4

结果预测也是4

可视化效果如下

七、自己手写的数字

由于自己手写的数字尺寸大小和通道与我们想要输入的图片不同，所以我们需要先对图片进行预处理。

新建一个py文件夹

预处理代码如下：

preprocess = transforms.Compose([
    transforms.Resize((28, 28))  # 尺寸大小
    , transforms.Grayscale(num_output_channels=1)  # 转为灰度图
    , transforms.ToTensor()  # 转换为张量
    , transforms.Normalize((0.1307,),(0.3081,))  # 归一化处理,这里的数据均值和方差是Mnist数据                
                                                    集均值和方差
])

整体代码加上可视化

"""
author: XiaoShu
date: 2024-08-21
"""
from torchvision import transforms
import train
import torch
from PIL import Image
import matplotlib.pyplot as plt

preprocess = transforms.Compose([
    transforms.Resize((28, 28))  # 尺寸大小
    , transforms.Grayscale(num_output_channels=1)  # 转为灰度图
    , transforms.ToTensor()  # 转换为张量
    , transforms.Normalize((0.1307,), (0.3081,))  # 归一化处理
])
# 定义路径
img_path1 = '0.png'
img_path2 = '3.png'
img_path3 = '5.png'
img_path4 = '8.png'
# 加载图片
image1 = Image.open(img_path1)
image2 = Image.open(img_path2)
image3 = Image.open(img_path3)
image4 = Image.open(img_path4)
# 预处理
image_1 = preprocess(image1)
image_2 = preprocess(image2)
image_3 = preprocess(image3)
image_4 = preprocess(image4)
# 修改成模型所需要的维度
image__1 = torch.reshape(image_1, (-1, 1, 28, 28))
image__2 = torch.reshape(image_2, (-1, 1, 28, 28))
image__3 = torch.reshape(image_3, (-1, 1, 28, 28))
image__4 = torch.reshape(image_4, (-1, 1, 28, 28))

model = train.Mnist()

model.load_state_dict(torch.load('Mnist1.pth', map_location=torch.device('cpu')))
model.eval()
with torch.no_grad():
    outPut1 = model(image__1)
    prediction1 = outPut1.argmax().item()
    outPut2 = model(image__2)
    prediction2 = outPut2.argmax().item()
    outPut3 = model(image__3)
    prediction3 = outPut3.argmax().item()
    outPut4 = model(image__4)
    prediction4 = outPut4.argmax().item()
    
    
# print(prediction1)
# print('识别结果：{}'.format(prediction1))

# 可视化
f, a = plt.subplots(2, 2, figsize=(7, 7))
a[0][0].set_title('predict:{}'.format(prediction1))
a[0][0].axis('off')
a[0][0].imshow(image1)

a[0][1].set_title('predict:{}'.format(prediction2))
a[0][1].axis('off')
a[0][1].imshow(image2)

a[1][0].set_title('predict:{}'.format(prediction3))
a[1][0].axis('off')
a[1][0].imshow(image3)

a[1][1].set_title('predict:{}'.format(prediction4))
a[1][1].axis('off')
a[1][1].imshow(image4)

plt.show()