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参数设置了每个包中的图片数据个数# 在装载的过程会将数据随机打乱顺序并进打包batch_size = 64#建立一个数据迭代器# 装载训练集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(train_loader))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 ~ 9def 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)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')LR = 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.0for i, data in enumerate(train_loader):inputs, labels = datainputs, 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 = 0total = 0for data_test in test_loader:images, labels = data_testimages, 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% 以上的成功率，效果还不错。
