实战：基于卷积的MNIST手写体分类

前面实现了基于多层感知机的MNIST手写体识别，本章将实现以卷积神经网络完成的MNIST手写体识别。

1. 数据的准备

在本例中，依旧使用MNIST数据集，对这个数据集的数据和标签介绍，前面的章节已详细说明过了，相对于前面章节直接对数据进行“折叠”处理，这里需要显式地标注出数据的通道，代码如下：

import numpy as npimport einops.layers.torch as elt#载入数据x_train = np.load("../dataset/mnist/x_train.npy")y_train_label = np.load("../dataset/mnist/y_train_label.npy")x_train = np.expand_dims(x_train,axis=1)   #在指定维度上进行扩充print(x_train.shape)

这里是对数据的修正，np.expand_dims的作用是在指定维度上进行扩充，这里在第二维（也就是PyTorch的通道维度）进行扩充，结果如下：

(60000, 1, 28, 28)

2. 模型的设计

下面使用PyTorch 2.0框架对模型进行设计，在本例中将使用卷积层对数据进行处理，完整的模型如下：

import torch
import torch.nn as nn
import numpy as np
import einops.layers.torch as elt
class MnistNetword(nn.Module):def __init__(self):super(MnistNetword, self).__init__()#前置的特征提取模块self.convs_stack = nn.Sequential(nn.Conv2d(1,12,kernel_size=7),  	#第一个卷积层nn.ReLU(),nn.Conv2d(12,24,kernel_size=5), 	#第二个卷积层nn.ReLU(),nn.Conv2d(24,6,kernel_size=3)  	#第三个卷积层)#最终分类器层self.logits_layer = nn.Linear(in_features=1536,out_features=10)def forward(self,inputs):image = inputsx = self.convs_stack(image)        #elt.Rearrange的作用是对输入数据的维度进行调整，读者可以使用torch.nn.Flatten函数完成此工作x = elt.Rearrange("b c h w -> b (c h w)")(x)logits = self.logits_layer(x)return logits
model = MnistNetword()
torch.save(model,"model.pth")

这里首先设定了3个卷积层作为前置的特征提取层，最后一个全连接层作为分类器层，需要注意的是，对于分类器的全连接层，输入维度需要手动计算，当然读者可以一步一步尝试打印特征提取层的结果，依次将结果作为下一层的输入维度。最后对模型进行保存。

3. 基于卷积的MNIST分类模型

下面进入本章的最后示例部分，也就是MNIST手写体的分类。完整的训练代码如下：

import torch
import torch.nn as nn
import numpy as np
import einops.layers.torch as elt
#载入数据
x_train = np.load("../dataset/mnist/x_train.npy")
y_train_label = np.load("../dataset/mnist/y_train_label.npy")
x_train = np.expand_dims(x_train,axis=1)
print(x_train.shape)
class MnistNetword(nn.Module):def __init__(self):super(MnistNetword, self).__init__()self.convs_stack = nn.Sequential(nn.Conv2d(1,12,kernel_size=7),nn.ReLU(),nn.Conv2d(12,24,kernel_size=5),nn.ReLU(),nn.Conv2d(24,6,kernel_size=3))self.logits_layer = nn.Linear(in_features=1536,out_features=10)def forward(self,inputs):image = inputsx = self.convs_stack(image)x = elt.Rearrange("b c h w -> b (c h w)")(x)logits = self.logits_layer(x)return logits
device = "cuda" if torch.cuda.is_available() else "cpu"
#注意记得将model发送到GPU计算
model = MnistNetword().to(device)
model = torch.compile(model)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
batch_size = 128
for epoch in range(42):train_num = len(x_train)//128train_loss = 0.for i in range(train_num):start = i * batch_sizeend = (i + 1) * batch_sizex_batch = torch.tensor(x_train[start:end]).to(device)y_batch = torch.tensor(y_train_label[start:end]).to(device)pred = model(x_batch)loss = loss_fn(pred, y_batch)optimizer.zero_grad()loss.backward()optimizer.step()train_loss += loss.item()  # 记录每个批次的损失值# 计算并打印损失值train_loss /= train_numaccuracy = (pred.argmax(1) == y_batch).type(torch.float32).sum().item() / batch_sizeprint("epoch：",epoch,"train_loss:", round(train_loss,2),"accuracy:",round(accuracy,2))

在这里，我们使用了本章新定义的卷积神经网络模块作为局部特征抽取，而对于其他的损失函数以及优化函数，只使用了与前期一样的模式进行模型训练。最终结果如下所示，请读者自行验证。

(60000, 1, 28, 28)
epoch： 0 train_loss: 2.3 accuracy: 0.11
epoch： 1 train_loss: 2.3 accuracy: 0.13
epoch： 2 train_loss: 2.3 accuracy: 0.2
epoch： 3 train_loss: 2.3 accuracy: 0.18
…
epoch： 58 train_loss: 0.5 accuracy: 0.98
epoch： 59 train_loss: 0.49 accuracy: 0.98
epoch： 60 train_loss: 0.49 accuracy: 0.98
epoch： 61 train_loss: 0.48 accuracy: 0.98
epoch： 62 train_loss: 0.48 accuracy: 0.98Process finished with exit code 0

本文节选自《PyTorch 2.0深度学习从零开始学》，本书实战案例丰富，可带领读者快速掌握深度学习算法及其常见案例。