#AI挑战营第一站# PyTorch实现MNIST手写数字识别

maskmoo

#AI挑战营第一站# PyTorch实现MNIST手写数字识别 [复制链接]

本文使用PyTorch框架，采用CNN卷积神经网络实现MNIST手写数字识别。

1 Pytorch环境安装

conda create -n pytorch python=3.6
conda activate pytorch
conda install pytorch torchvision torchaudio cpuonly -c pytorch

 

  如遇到网络引起的安装问题可以下来离线安装

安装pytorch(cpu版)，torch和torchvision的详细安装步骤

 

Pytorch环境检查

print(torch.__version__)输出的时torch的版本，print(torch.cuda.is_available())的结果如果是False代表cpu版本。

 

2 数据集（MNIST）

MNIST数据集是机器学习领域中非常经典的一个数据集，可以通过官方网址进行下载http://yann.lecun.com/exdb/mnist/ 。或者后面数据加载时加入download=True选项。

 

1. 数据预处理： 使用了transforms.Compose来定义数据预处理管道，其中包括将图像转换为张量和归一化处理，使得图像数据的像素值在0到1之间。

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

2. 数据加载： 使用了torchvision.datasets中的MNIST数据集，通过DataLoader加载数据，并设置了批量大小和是否打乱数据的选项。

train_dataset = datasets.MNIST(root='./data/', train=True, download=True, transform=transform)  
test_dataset = datasets.MNIST(root='./data/', train=False, download=True, transform=transform)  # train=True训练集，=False测试集

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

3 构建模型（CNN）

1. 模型结构： 通过定义一个继承自torch.nn.Module的类来构建CNN模型，其中包括卷积层、激活函数和池化层。

 

 

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10),
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)  
        x = self.conv2(x) 
        x = x.view(batch_size, -1) 
        x = self.fc(x)
        return x 

1. 模型训练： 在模型类中定义了forward方法，用于前向传播，同时定义了损失函数（交叉熵损失）和优化器（随机梯度下降）。

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum) 

   定义训练轮和测试轮

1. 训练轮： 定义了一个训练函数，其中包括前向传播、计算损失、反向传播和参数更新的步骤。

   def train(epoch):
       running_loss = 0.0
       running_total = 0
       running_correct = 0
       for batch_idx, data in enumerate(train_loader, 0):
           inputs, target = data
           optimizer.zero_grad()
   
           # forward + backward + update
           outputs = model(inputs)
           loss = criterion(outputs, target)
   
           loss.backward()
           optimizer.step()
   
           # 把运行中的loss累加起来，为了下面300次一除
           running_loss += loss.item()
           # 把运行中的准确率acc算出来
           _, predicted = torch.max(outputs.data, dim=1)
           running_total += inputs.shape[0]
           running_correct += (predicted == target).sum().item()
   
           if batch_idx % 300 == 299:  # 不想要每一次都出loss，浪费时间，选择每300次出一个平均损失,和准确率
               print('[%d, %5d]: loss: %.3f , acc: %.2f %%'
                     % (epoch + 1, batch_idx + 1, running_loss / 300, 100 * running_correct / running_total))
               running_loss = 0.0  # 这小批300的loss清零
               running_total = 0
               running_correct = 0  # 这小批300的acc清零
   

    

2. 测试轮： 定义了一个测试函数，用于评估模型在测试集上的性能。

   def test():
       correct = 0
       total = 0
       with torch.no_grad():  # 测试集不用算梯度
           for data in test_loader:
               images, labels = data
               outputs = model(images)
               _, predicted = torch.max(outputs.data, dim=1)  # dim = 1 列是第0个维度，行是第1个维度，沿着行(第1个维度)去找1.最大值和2.最大值的下标
               total += labels.size(0)  # 张量之间的比较运算
               correct += (predicted == labels).sum().item()
       acc = correct / total
       print('[%d / %d]: Accuracy on test set: %.1f %% ' % (epoch+1, EPOCH, 100 * acc))  # 求测试的准确率，正确数/总数
       return acc
   

    

开始训练

1. 超参数设置： 设置了批量大小、学习率和冲量等超参数，这些超参数的选择对模型的训练效果有重要影响  
2. 训练过程： 使用一个循环迭代多个轮次进行训练，每轮训练后进行一次测试。

保存参数

        torch.save(model.state_dict(), './model_mnist.pth')
        torch.save(optimizer.state_dict(), './optimizer_mnist.pth')

model.state_dict()和optimizer.state_dict()都是PyTorch中用于保存和加载模型状态的函数，但它们保存的内容不同：

1. model.state_dict():

   • model.state_dict()返回一个Python字典，其中包含了模型的所有参数（权重和偏置），以及各个层的名称和参数对应的张量值。
   • 通常用于保存和加载模型的参数状态。可以将这个字典保存到文件中，并在需要时使用load_state_dict()方法加载到模型中。

2. optimizer.state_dict():

   • optimizer.state_dict()返回一个Python字典，其中包含了优化器的状态信息，包括当前参数和动量缓存。
   • 通常用于保存和加载优化器的状态，以便在训练中断后恢复训练。

通常会同时保存模型的状态和优化器的状态，以便在需要时能够完全恢复模型的状态和训练状态，而无需重新开始训练。

4 ONNX转换

1. 定义模型: 首先，确保已经定义了要导出的PyTorch模型，并加载了其参数。

2. 准备输入: 准备一个示例输入用于在导出过程中推断模型的输入形状。

3. 导出模型: 使用torch.onnx.export()函数导出模型。需要指定模型、示例输入、导出文件路径以及其他参数。

    #导出为onnx模型
    input = torch.randn(1, 28, 28)

    torch.onnx.export(model, input, "mnist.onnx", verbose=True)

ONNX查看

 

Netron是一个轻量级的 ONNX 模型查看器，可以在浏览器中运行。通过使用 pip 安装 Netron：

pip install netron

然后，您可以使用以下命令启动 Netron 并打开 ONNX 文件：

netron <your_model.onnx>

 

 

5 完整代码

 

以下代码用于训练 MNIST 手写数字识别模型，并在训练过程中输出训练轮次的损失和准确率。程序包括了以下主要部分

1. 导入必要的库：导入了 PyTorch、NumPy、Matplotlib 和 DataLoader 等库，以及 MNIST 数据集和一些转换操作。

2. 超参数设置：定义了批量大小、学习率、动量和训练轮次等超参数。

3. 数据准备：定义了 MNIST 数据集的转换和加载器。

4. 模型设计：使用 torch.nn.Sequential 构建了一个简单的卷积神经网络模型，其中包括两个卷积层和两个全连接层。

5. 损失函数和优化器：定义了交叉熵损失函数和随机梯度下降（SGD）优化器。

6. 训练和测试函数：定义了 train 和 test 函数，用于训练和测试模型。在训练过程中，通过迭代数据加载器中的数据批次，计算损失、更新模型参数，并输出训练轮次的损失和准确率。在测试过程中，通过迭代测试数据加载器中的数据批次，评估模型在测试集上的准确率。

7. 命令行参数解析：使用 argparse 模块解析命令行参数，以便在命令行中指定训练轮次。

8. 训练和测试过程：在 if __name__ == '__main__': 块中，加载命令行参数并开始训练和测试模型。在每个训练轮次结束后，检查测试准确率是否达到最佳，如果是，则保存模型参数。最后，绘制测试准确率随训练轮次变化的曲线。

9. 导出 ONNX 模型：在训练完成后，加载表现最佳的模型参数，并导出为 ONNX 格式的模型文件。

import torch
import numpy as np
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import torch.nn.functional as F
import argparse

# OMP: Error #15
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

"""
卷积运算 使用mnist数据集，和10-4，11类似的，只是这里：1.输出训练轮的acc 2.模型上使用torch.nn.Sequential
"""
# Super parameter ------------------------------------------------------------------------------------
batch_size = 64
learning_rate = 0.01
momentum = 0.5
EPOCH = 10

# Prepare dataset ------------------------------------------------------------------------------------
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
# softmax归一化指数函数(https://blog.csdn.net/lz_peter/article/details/84574716),其中0.1307是mean均值和0.3081是std标准差

train_dataset = datasets.MNIST(root='./data/mnist', train=True, transform=transform, download=True)  # 本地没有就加上download=True
test_dataset = datasets.MNIST(root='./data/mnist', train=False, transform=transform)  # train=True训练集，=False测试集
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# fig = plt.figure()
# for i in range(12):
#     plt.subplot(3, 4, i+1)
#     plt.tight_layout()
#     plt.imshow(train_dataset.train_data, cmap='gray', interpolation='none')
#     plt.title("Labels: {}".format(train_dataset.train_labels))
#     plt.xticks([])
#     plt.yticks([])
# plt.show()


# 训练集乱序，测试集有序
# Design model using class ------------------------------------------------------------------------------
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10),
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)  # 一层卷积层,一层池化层,一层激活层(图是先卷积后激活再池化，差别不大)
        x = self.conv2(x)  # 再来一次
        x = x.view(batch_size, -1)  # flatten 变成全连接网络需要的输入 (batch, 20,4,4) ==> (batch,320), -1 此处自动算出的是320
        x = self.fc(x)
        return x  # 最后输出的是维度为10的，也就是（对应数学符号的0~9）


model = Net()


# Construct loss and optimizer ------------------------------------------------------------------------------
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)  # lr学习率，momentum冲量




# Train and Test CLASS --------------------------------------------------------------------------------------
# 把单独的一轮一环封装在函数类里
def train(epoch):
    running_loss = 0.0  # 这整个epoch的loss清零
    running_total = 0
    running_correct = 0
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, target = data
        optimizer.zero_grad()

        # forward + backward + update
        outputs = model(inputs)
        loss = criterion(outputs, target)

        loss.backward()
        optimizer.step()

        # 把运行中的loss累加起来，为了下面300次一除
        running_loss += loss.item()
        # 把运行中的准确率acc算出来
        _, predicted = torch.max(outputs.data, dim=1)
        running_total += inputs.shape[0]
        running_correct += (predicted == target).sum().item()

        if batch_idx % 300 == 299:  # 不想要每一次都出loss，浪费时间，选择每300次出一个平均损失,和准确率
            print('[%d, %5d]: loss: %.3f , acc: %.2f %%'
                  % (epoch + 1, batch_idx + 1, running_loss / 300, 100 * running_correct / running_total))
            running_loss = 0.0  # 这小批300的loss清零
            running_total = 0
            running_correct = 0  # 这小批300的acc清零

        # torch.save(optimizer.state_dict(), './optimizer_mnist.pth')


def test():
    correct = 0
    total = 0
    with torch.no_grad():  # 测试集不用算梯度
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, dim=1)  # dim = 1 列是第0个维度，行是第1个维度，沿着行(第1个维度)去找1.最大值和2.最大值的下标
            total += labels.size(0)  # 张量之间的比较运算
            correct += (predicted == labels).sum().item()
    acc = correct / total
    print('[%d / %d]: Accuracy on test set: %.1f %% ' % (epoch+1, EPOCH, 100 * acc))  # 求测试的准确率，正确数/总数
    return acc


# Start train and Test --------------------------------------------------------------------------------------

if __name__ == '__main__':
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='Train a model with specified number of epochs.')

    # 添加 --epochs 参数
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')

    # 解析命令行参数
    args = parser.parse_args()

    # 使用 args.epochs 获取命令行传入的值
    EPOCH = epochs = args.epochs

    # 打印参数值
    print(f'Number of epochs to train: {epochs}')
    acc_list_test = []
    acc_best = 0
    epoch_bset = 0
    for epoch in range(EPOCH):
        train(epoch)
        # if epoch % 10 == 9:  #每训练10轮 测试1次
        acc_test = test()
        if acc_test > acc_best:
            acc_best = acc_test
            epoch_bset = epoch
            torch.save(model.state_dict(), './model_mnist_best.pth')
        acc_list_test.append(acc_test)
    print('Bset:[%d / %d]: Accuracy on test set: %.1f %% ' % (epoch_bset+1, EPOCH, 100 * acc_best))  # 求测试的准确率，正确数/总数

    plt.plot(acc_list_test)
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy On TestSet')
    plt.show()
    
    #导出为onnx模型
    model.load_state_dict(torch.load('model_mnist_best.pth'))
    input = torch.randn(1, 28, 28)
    torch.onnx.export(model, input, "mnist.onnx", verbose=True)

 

 

6 参考链接

1. 全面总结 pip install 与 conda install 的使用区别_pytorch安装pip和conda的区别
2. (详解)安装pytorch(cpu版)，torch和torchvision的详细安装步骤，并在jupyter notebook和pycharm中配置好pytorch_怎么安装匹配的torchvision
3. 用PyTorch实现MNIST手写数字识别_mnist pytorch