基于Pytorch的Mnist手写数字识别

# !/usr/bin/env Python3
# -*- coding: utf-8 -*-
# @version: v1.0
# @Author   : Meng Li
# @contact: 925762221@qq.com
# @FILE     : torch_mnist.py
# @Time     : 2022/5/31 9:29
# @Software : PyCharm
# @site:
# @Description : 自己动手实现mnist数据集的10分类任务
# 同等条件下，batch_size 越小，模型越收敛。但是更容易震荡。learning_rate越小，模型收敛速度越慢

import torch
import torch.nn as nn
import torchvision
import torch.nn.functional as F
import torchsummary
import torch.optim as optim
from torch.utils.data import Dataset
import matplotlib.pylab as plt


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, 100)
        self.fc2 = nn.Linear(100, 10)
        self.crition = torch.nn.CrossEntropyLoss()
        pass

    def forward(self, x, y):
        batch_size, _, h, w = x.size()
        x = x.view(-1, h * w)
        output = F.relu(self.fc1(x))
        output = self.fc2(output)
        loss = self.crition(output, y)
        val, index = torch.max(output, 1)
        acc = torch.eq(index, y).float().cpu().sum()
        return loss, acc.float() / y.size(0), index


def train():
    net = Net()
    show_sum_flg = False
    if show_sum_flg:
        torchsummary.summary(net, (28, 28))
    train_data = torchvision.datasets.MNIST(root="./", train=True, transform=torchvision.transforms.ToTensor(),
                                            download=False)
    batch_size = 64
    learning_rate = 0.001
    optimizer = optim.SGD(net.parameters(), lr=learning_rate)

    train_iter = torch.utils.data.DataLoader(train_data, batch_size, shuffle=True)
    epoch = 30
    max_acc = 0
    acc = 0
    for i in range(epoch):
        for image, label in train_iter:
            loss, acc, _ = net(image, label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print("epoch {0}  acc {1}".format(i, acc))
        if acc > max_acc:
            max_acc = acc
            torch.save(net, 'limeng.pth')


def test():
    net = torch.load('limeng.pth')
    net.eval()
    train_data = torchvision.datasets.MNIST(root="./", train=True, transform=torchvision.transforms.ToTensor(),
                                            download=False)
    batch_size = 10
    train_iter = torch.utils.data.DataLoader(train_data, batch_size, shuffle=True)

    for image, label in train_iter:
        _, _, predict = net(image, label)
        for i in range(batch_size):
            imagei = image[i, 0, :, :]
            plt.subplot(2, 5, i+1)
            plt.imshow(imagei)
            plt.title("{0}".format(predict[i]))
        plt.show()
        break


if __name__ == '__main__':
    # train()
    test()

先上代码，工作期间接触了Tensorflow和Pytorch两种框架，但是总得来说,pytorch由于编码语法规范更接近于python原生语法，所以更容易上手。作为深度学习中的"hello world"，还是有必要自己写一下整个数据输入到模型训练，模型保存再到模型测试的全流程。

测试模型时，运行效果图大概是这样的：