基于pytorch使用MNIST数据集进行图像识别

图像识别是人工智能中的一个基础任务，它的目标是让计算机能够识别图像中的物体等内容，并将它们分配到预定义的类别中。例如，给定一张手写数字的图片，图像识别系统应该能够输出这个图片对应的值的类别。

为了训练和评估图像识别系统，需要有大量的带有标注的图像数据集。常用的图像分类数据集有：

• ImageNet：一个包含超过1400万张图片和2万多个类别的大型数据库，是目前最流行和最具挑战性的图像分类基准之一。
• CIFAR-10/CIFAR-100：一个包含6万张32×32大小的彩色图片和10或100个类别的小型数据库，适合入门级和快速实验。
• MNIST：一个包含7万张28×28大小的灰度手写数字图片和10个类别的经典数据库，是深度学习中最常用的测试集之一。

本文将讲述如何使用上述MNIST数据集进行图像识别任务。后续有机会会讲述图像卷积模型以及如何使用图像对抗来干扰识别任务。

首先导入依赖库

import torch
import torchvision
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

使用torchvision库可以方便地加载这些常用数据集或者自定义数据集。torchvision.datasets提供了一些加载数据集或者下载数据集到本地缓存文件夹（默认为./data）并返回Dataset对象（torch.utils.data.Dataset）的函数。

例如，要加载MNIST训练集，首先需要下载并加载数据

trans_to_tensor = transforms.Compose([transforms.ToTensor()]) # 定义转换函数，将PIL.Image转换为torch.Tensor

torchvision.transforms.ToTensor这个函数能把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray，转换为形状为[C,H,W]，取值范围是[0,1.0]的torch.FloatTensor。其中(H,W,C)意思是(高，宽，通道数)，黑白图片的通道数只有1，其中每个像素点的取值为[0,255]，彩色图片的通道数为(R,G,B),每个通道的像素点的取值为[0,255]，三个通道的颜色相互叠加，形成了各种颜色。

torchvision.transforms.Compose(transforms)可以向其中传入一个列表，列表中放入各种transforms的对象，例如

transforms.Compose([
 torchvision.transforms.ToTensor()
 torchvision.transforms.Normalize(mean,std)
])

上面的代码可以先转换为Tensor，再进行正则化等操作

# 从 torchvision 中加载 MNIST 数据集的训练集
# - root='./data': 数据集将被下载并保存在当前工作目录下的 'data' 子目录中
# - train=True: 加载训练集，false加载测试集
# - download=True: 如果数据集不存在，则下载数据集
# - transform=transform: 对加载的图像应用上面定义的图像变换操作
data_train = torchvision.datasets.MNIST(root='./data', train=True, transform=trans_to_tensor, download=True)
data_test = torchvision.datasets.MNIST(root='./data', train=False, transform=trans_to_tensor, download=True)

定义数据加载器

# 创建一个 DataLoader 对象，用于对数据进行批量加载和处理
# - data_train: 要加载的数据集
# - batch_size=100: 每个批次包含的图像样本数量
# - shuffle=True: 打乱数据，以便在每个 epoch 中随机访问样本
train_loader = torch.utils.data.DataLoader(data_train, batch_size=100, shuffle=True) # 定义DataLoader对象，用于批量加载数据

以上代码的关键作用：

• 定义了图像的预处理操作，将原始图像数据转换成深度学习模型可以处理的格式。
• 下载并加载了 MNIST 数据集的训练集，同时应用了已经定义的图像预处理操作。
• 创建了一个数据加载器，用于在训练过程中按照批次加载和处理图像数据。

接下来就是定义一个神经网络模型，通过继承torch.nn.Module类来实现一个Module对象，并实现init和forward两个方法。

class Net(nn.Module):
    # init方法用于定义模型中需要的各种层和参数
    def __init__(self):
        super(Net, self).__init__()
        # 定义神经网络的全连接层（线性层）,全连接层是指当前一层的神经元和前一层的神经元相互链接，其核心操作就是y=wx，即矩阵的乘法，实现对前一层数据的变换。
        # 输入层
        self.fc1 = nn.Linear(1*28*28,28)  
        # 输出层（10个类别）
        self.fc2 = nn.Linear(28,10)        
    
    # forward方法用于定义前向传播过程，即如何根据输入的图像张量（Tensor）计算出输出的类别概率分布。我们可以使用定义好的各种层和参数，并结合一些激活函数（如ReLU）和归一化函数（如softmax）来实现forward方法。   
    def forward(self, x):
        # 前向传播函数，定义了数据从输入到输出的流程
        # 进行形状的修改,将输入的图像数据展平成一维向量
        x = x.view(-1, 1 * 28 * 28)   
        # 进行全连接的操作         
        x = self.fc1(x)        
        # ReLU是一种激活函数，即整流线性单元（Rectified Linear Unit），用于神经网络中的前向传播过程。它将所有负数值的输出都设为零，而正数值的输出则不做改变。这个函数可以帮助神经网络更快地收敛，提高运算速度。激活函数的处理过程中,形状没有变化。
        x = F.relu(x)       
        # 输出层   
        out = self.fc2(x)                    
        return out

创建一个网络实例

net = Net()

定义损失函数，作为损失函数的意义是：当预测结果越接近真实值，损失函数的值越接近于0

# 交叉熵损失函数，用于多类别分类任务。
criterion = nn.CrossEntropyLoss()  

定义优化器，优化器主要是在模型训练阶段对模型可学习参数进行更新,优化算法的功能，是通过改善训练方式，来最小化（或最大化）损失函数。

# 随机梯度下降优化器
optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9) 

优化器在训练过程中，需要：

• 先调用 optimizer.zero_grad() 清空梯度
• 再调用 loss.backward() 反向传播
• 最后调用 optimizer.step()更新模型参数

模型训练模块构建

# 在训练模型时加上 net.train()，这样会正常使用 Batch Normalization 和 Dropout
net.train() # 将模型(包含所有子模块)中的参数转换成训练状态
for epoch in range(5):  # 循环训练5个周期
    running_loss = 0.0 
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        # 梯度清零，防止梯度累积
        optimizer.zero_grad()         
        # 前向传播，计算网络输出
        outputs = net(inputs)         
        # 将预测值和真实值传入，计算损失
        loss = criterion(outputs, labels)  
        # 反向传播，计算梯度
        loss.backward()  
        # 更新网络参数             
        optimizer.step()
        # 累积损失值              
        running_loss += loss.item()   
        # 批次打印一次损失值
        if i % 100 == 99:  
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')
            running_loss = 0.0

运行输出

[1,   100] loss: 2.233
[1,   200] loss: 2.028
[1,   300] loss: 1.758
[1,   400] loss: 1.445
[1,   500] loss: 1.190
[1,   600] loss: 0.983
[2,   100] loss: 0.847
[2,   200] loss: 0.750
[2,   300] loss: 0.667
[2,   400] loss: 0.619
[2,   500] loss: 0.585
[2,   600] loss: 0.542
[3,   100] loss: 0.527
[3,   200] loss: 0.500
[3,   300] loss: 0.495
[3,   400] loss: 0.459
[3,   500] loss: 0.449
[3,   600] loss: 0.436
[4,   100] loss: 0.432
[4,   200] loss: 0.420
[4,   300] loss: 0.414
[4,   400] loss: 0.411
[4,   500] loss: 0.400
[4,   600] loss: 0.392
[5,   100] loss: 0.394
[5,   200] loss: 0.386
[5,   300] loss: 0.382
[5,   400] loss: 0.374
[5,   500] loss: 0.362
[5,   600] loss: 0.364

# 测试的时候一般选择 net.eval()，这样就不会使用 Batch Normalization 和 Dropout
net.eval() # 将模型(包含所有子模块)中的参数转换成预测状态
test_loader = torch.utils.data.DataLoader(data_test, batch_size=10000, shuffle=False)
test_data = next(iter(test_loader)) # 迭代器对象从集合的第一个元素开始访问，直到所有的元素被访问完结束。迭代器只能往前不会后退。
with torch.no_grad():
    x, y = test_data[0], test_data[1]
    outputs = net(x)
    pred = torch.max(outputs, 1)[1]
    print(f'test acc: {sum(pred == y)} / {y.shape[0]}')

对测试集进行测试的结果为：test acc: 9002 / 10000

后续可以对单个图像进行测试

# 获取一个图像
data_demo = data_train[1][0]
# transpose(1,2,0) 将(1, 28, 28)的shape转换为(28, 28, 1)
origin= data_demo.numpy().transpose(1,2,0) 
plt.imshow(origin,cmap='Greys')

# 对该图像进行分类预测
output = net(data_demo)
torch.max(output,1)[1].item()

显示结果为：0

# 保存模型
torch.save(net, 'nn.pt')

最后就是模型的保存，以便下次直接调用。