PyToch：基于神经网络的数字识别（MNIST数据集）

背景

最近在学习PyTorch和深度学习，所以决定先用MNIST数据集试试手，利用神经网络做一个简单的数字识别。
参考代码来源于：
https://github.com/udacity/deep-learning-v2-pytorch/tree/master/intro-to-pytorch
需要注意的是本文的神经网络中只使用了全链接层，只是为了理解PyTorch而进行的一个非常简单的实验。
以下的代码内容包括train.py和predict.py，后边就直接边贴代码边解释。

代码链接：https://github.com/Yannnnnnnnnnnn/learnPyTorch/tree/master/trainMNIST/fulljion

---

一、训练

train_CPU.py

# 加载训练需要的模块
import torch
from torch import nn
from torch import optim
from torchvision import datasets, transforms

# transform模块，主要作用是将输入数据转换成tenssor，并进行归一化
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,)),
                              ])

# MNIST数据 train,每一个batch的大小为128
trainset = datasets.MNIST('./MNIST_data/', download=True, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

# MNIST数据 test,每一个batch的大小为128
testset = datasets.MNIST('./MNIST_data/', download=True, train=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=True)


# 本次实验所使用的模型，只有全链接层，结构非常简单，没有什么特殊的地方
model = nn.Sequential(nn.Linear(784, 392),
                      nn.ReLU(),
                      nn.Linear(392, 128),
                      nn.ReLU(),
                      nn.Linear(128, 64),
                      nn.ReLU(),
                      nn.Linear(64, 10),
                      nn.LogSoftmax(dim=1))

# 定义一下损失函数和优化器
criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.003)

# 迭代次数
epochs = 30

# 记录每次的损失
train_losses, test_losses = [], []

# 训练
for e in range(epochs):
    running_loss = 0
    # 读取所有的训练数据，并进行训练
    # images和labels都是维度为128的tensor
    for images, labels in trainloader:
        
        # 将images的size更改一下，变成一条向量
        images = images.view(images.shape[0], -1)
    
        # 清除上一次自动求导的梯度信息
        optimizer.zero_grad()
        
        # forward过程
        output = model(images)
        
        # 计算损失
        loss = criterion(output, labels)
        
        # backward 
        # 此过程中会自动求导
        loss.backward()
        
        # 更新参数
        optimizer.step()
        
        running_loss += loss.item()
    else:
        test_loss = 0
        accuracy = 0
        
        # 利用test数据进行测试
        # 为提高预测的速度，最好关闭梯度计算
        with torch.no_grad():
            for images, labels in testloader:
                
                # 更改images的size
                images = images.view(images.shape[0], -1)
                
                # 预测
                log_ps = model(images)
                
                # 计算损失
                test_loss += criterion(log_ps, labels)
                
                # 由于在pytorch中最终的预测结果都进行了求对数
                # 所以这这里又添加了一个求指数
                ps = torch.exp(log_ps)
                
                # 获取最好的结果
                top_p, top_class = ps.topk(1, dim=1)
                
                # 计算精度
                equals = top_class == labels.view(*top_class.shape)
                accuracy += torch.mean(equals.type(torch.FloatTensor))
                
        train_losses.append(running_loss/len(trainloader))
        test_losses.append(test_loss/len(testloader))

        # 在每一个epoch中都保存一次模型
        torch.save(model.state_dict(), str(e) +'.pth')

        print("Epoch: {}/{}.. ".format( e+1, epochs),
              "Training Loss: {:.3f}.. ".format(running_loss/len(trainloader)),
              "Test Loss: {:.3f}.. ".format(test_loss/len(testloader)),
              "Test Accuracy: {:.3f}".format(accuracy/len(testloader)))
      
# 画一下最终的精度图        
import matplotlib.pyplot as plt
plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)

以上代码都是在CPU上运行，如果要更改成gpu也非常简单，以下仅贴出了需要修改的区域

# 本次实验所使用的模型，只有全链接层，结构非常简单，没有什么特殊的地方
model = nn.Sequential(nn.Linear(784, 392),
                      nn.ReLU(),
                      nn.Linear(392, 128),
                      nn.ReLU(),
                      nn.Linear(128, 64),
                      nn.ReLU(),
                      nn.Linear(64, 10),
                      nn.LogSoftmax(dim=1))
model.cuda() # GPU

# 训练
        # 将images的size更改一下，变成一条向量
        images = images.view(images.shape[0], -1)
        
        # CUDA
        images = images.cuda()
        labels = labels.cuda()
    
 
    else:

     
        # 利用test数据进行测试
        # 为提高预测的速度，最好关闭梯度计算
        with torch.no_grad():
            for images, labels in testloader:
                
                # 更改images的size
                images = images.view(images.shape[0], -1)
                
                # CUDA
                images = images.cuda()
                labels = labels.cuda()
    

利用以上代码就可以进行训练，最后经过30次迭代后迭代情况如下图所示，最后的validation精度为92.8%

---

二、预测

完成训练后，我分别用MNIST中的几个数字和自己用画图板写了几个数字进行测试。
MNIST的数据如下：

自己画的数据如下：

以下是预测所使用的代码
predict.py

import torch
from torch import nn
from torchvision import transforms

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


model = nn.Sequential(nn.Linear(784, 392),
                      nn.ReLU(),
                      nn.Linear(392, 128),
                      nn.ReLU(),
                      nn.Linear(128, 64),
                      nn.ReLU(),
                      nn.Linear(64, 10),
                      nn.LogSoftmax(dim=1))



# load model
state_dict = torch.load('29.pth')
model.load_state_dict(state_dict)


import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# 打开图片并转成灰度
image = Image.open('1.png')
gray=image.convert('L')

# 显示图片
plt.figure("predict")
plt.imshow(gray)
plt.show()


# 转成tensor
tensor = transform(gray)
tensor = tensor.view(1, 784)
inputdata = torch.autograd.Variable(tensor,requires_grad=False)
outputdata = model(inputdata)
ps = torch.exp(outputdata)

top_p, top_class = ps.topk(1, dim=1)

# 输出结果
print(top_p)

最后在MNIST数据上实验结果如下：

0(0.9987) 1(0.9855) 2(0.9785) 5(0.7746) 9(0.9116)

在自己画的数据上的实验结果如下，其中第4个数字7被错误分成了2。想来全链接还是太简单了

1(0.5416) 3(0.4719) 4(0.6279) 2(0.5658) 8(0.7286)

---

三、总结

本博客只是简单的尝试了一下，并且只使用了全链接，所以结果并不是很好。以后尝试加入卷积层，效果应该会更好。