MNIST 手写数字识别是一个比较简单的入门项目，相当于深度学习中的 Hello World，可以让我们快速了解构建神经网络的大致过程。虽然网上的案例比较多，但还是要自己实现一遍。代码采用 PyTorch 1.0 编写并运行。

导入相关库

import torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optimfrom torchvision import datasets, transformsimport torchvisionfrom torch.autograd import Variablefrom torch.utils.data import DataLoaderimport cv2

torchvision 用于下载并导入数据集

cv2 用于展示数据的图像

获取训练集和测试集

# 下载训练集train_dataset = datasets.MNIST(root='./num/', train=True, transform=transforms.ToTensor(), download=True)# 下载测试集test_dataset = datasets.MNIST(root='./num/', train=False, transform=transforms.ToTensor(), download=True)

root 用于指定数据集在下载之后的存放路径

transform 用于指定导入数据集需要对数据进行那种变化操作

train是指定在数据集下载完成后需要载入的那部分数据，设置为 True 则说明载入的是该数据集的训练集部分，设置为 False 则说明载入的是该数据集的测试集部分

download 为 True 表示数据集需要程序自动帮你下载

这样设置并运行后，就会在指定路径中下载 MNIST 数据集，之后就可以使用了。

数据装载和预览

# dataset 参数用于指定我们载入的数据集名称# batch_size参数设置了每个包中的图片数据个数# 在装载的过程会将数据随机打乱顺序并进打包# 装载训练集train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)# 装载测试集test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

在装载完成后，可以选取其中一个批次的数据进行预览：

images, labels = next(iter(data_loader_train))img = torchvision.utils.make_grid(images)img = img.numpy().transpose(1, 2, 0)std = [0.5, 0.5, 0.5]mean = [0.5, 0.5, 0.5]img = img * std + meanprint(labels)cv2.imshow('win', img)key_pressed = cv2.waitKey(0)

在以上代码中使用了 iter 和 next 来获取取一个批次的图片数据和其对应的图片标签，然后使用 torchvision.utils 中的 make_grid 类方法将一个批次的图片构造成网格模式。

预览图片如下：

并且打印出了图片相对应的数字：

搭建神经网络

# 卷积层使用 torch.nn.Conv2d# 激活层使用 torch.nn.ReLU# 池化层使用 torch.nn.MaxPool2d# 全连接层使用 torch.nn.Linearclass LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.conv1 = nn.Sequential(nn.Conv2d(1, 6, 3, 1, 2), nn.ReLU(), nn.MaxPool2d(2, 2)) self.conv2 = nn.Sequential(nn.Conv2d(6, 16, 5), nn.ReLU(), nn.MaxPool2d(2, 2)) self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120), nn.BatchNorm1d(120), nn.ReLU()) self.fc2 = nn.Sequential( nn.Linear(120, 84), nn.BatchNorm1d(84), nn.ReLU(), nn.Linear(84, 10)) # 最后的结果一定要变为 10，因为数字的选项是 0 ~ 9 def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = x.view(x.size()[0], -1) x = self.fc1(x) x = self.fc2(x) x = self.fc3(x) return x

前向传播内容：

首先经过 self.conv1() 和 self.conv1() 进行卷积处理

然后进行 x = x.view(x.size()[0], -1)，对参数实现扁平化（便于后面全连接层输入）

最后通过 self.fc1() 和 self.fc2() 定义的全连接层进行最后的分类

训练模型

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')batch_size = 64LR = 0.001net = LeNet().to(device)# 损失函数使用交叉熵criterion = nn.CrossEntropyLoss()# 优化函数使用 Adam 自适应优化算法optimizer = optim.Adam( net.parameters(), lr=LR,)epoch = 1if __name__ == '__main__': for epoch in range(epoch): sum_loss = 0.0 for i, data in enumerate(train_loader): inputs, labels = data inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda() optimizer.zero_grad() #将梯度归零 outputs = net(inputs) #将数据传入网络进行前向运算 loss = criterion(outputs, labels) #得到损失函数 loss.backward() #反向传播 optimizer.step() #通过梯度做一步参数更新 # print(loss) sum_loss += loss.item() if i % 100 == 99: print('[%d,%d] loss:%.03f' % (epoch + 1, i + 1, sum_loss / 100)) sum_loss = 0.0

测试模型

net.eval() #将模型变换为测试模式 correct = 0 total = 0 for data_test in test_loader: images, labels = data_test images, labels = Variable(images).cuda(), Variable(labels).cuda() output_test = net(images) _, predicted = torch.max(output_test, 1) total += labels.size(0) correct += (predicted == labels).sum() print("correct1: ", correct) print("Test acc: {0}".format(correct.item() / len(test_dataset)))

训练及测试的情况：

98% 以上的成功率，效果还不错。

以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持。