MNIST手写数字识别是pytorch的经典项目之一，适用于新手学习pytorch这一深度学习训练平台以及初步了解计算机视觉的相关训练步骤。

在新手学习pytorch的时候，最好的教程就是Pytorch official Tutorial。以下的网络搭建主要参考pytorch的官方教程完成。

1
2
3
4
5
6
7
8
import torch 
import torchvision
import torchvision.transforms as transforms

# 在windows显示图像的时候可能会遇到一个报错，需要添加这个语句才可以正常通过，意思是允许重复加载动态链接库
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

首先引入几个重要的Package。

1
2
3
4
5
6
7

device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
dataset_root = './data'
batch_size  = 64
epoch_time = 20
path = './checkpoints/MNIST.pth'

接下来需要对几个重要的全局参数进行设置，device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')用来设置device为CPU还是GPU，由于GPU进行矩阵运算更快因此尽可能将网络搭载在GPU上。
之后的几个参数也主要是设置数据集存储位置和batch_size和训练轮数epoch_time以及训练网络参数的存储位置path。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

transform = transforms.Compose([
    transforms.ToTensor(),

    transforms.Normalize((0.5,),(0.5,))
])

trainset = torchvision.datasets.MNIST(root=dataset_root,train = True,transform=transform,download=True)

trainloader = torch.utils.data.DataLoader(trainset,batch_size=batch_size,num_workers=0,shuffle=True)

testset = torchvision.datasets.MNIST(root=dataset_root,train= False,download=True,transform=transform)

testloader = torch.utils.data.DataLoader(testset,batch_size=batch_size,num_workers=0,shuffle=False)

这部分代码主要实现数据集的下载和分装。分别加载到对应的数据加载器，方便后续的网络训练。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
'''
show image test
'''
import matplotlib.pyplot as plt 
import numpy as np 


def imshow(img):
    npimg = img.numpy()
    plt.imshow(npimg.transpose(1,2,0))
    plt.show()


dataiter = iter(trainloader)
# 注意这里不能够使用*.to(device)将数据放到GPU上，因为是要进行图像展示，处理需要在CPU上进行
images,labels = dataiter.next()

imshow(torchvision.utils.make_grid(images))
print(labels)

以上部分代码主要实现部分图像的显示。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

'''
    design a module 
'''
import torch.nn as nn
import torch.nn.functional as F  


class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        # 1*1*28*28
        self.conv1 = nn.Conv2d(1,6,5)
        self.conv2 = nn.Conv2d(6,16,5)
        self.fc1 = nn.Linear(16*8*8,120)
        self.fc2 = nn.Linear(120,10)
        self.pool = nn.MaxPool2d(2,2)

    def forward(self,x): 
        x= self.pool(F.relu(self.conv1(x)))   
        x = F.relu(self.conv2(x))
        x = x.view(-1,16*8*8)
        x= F.relu(self.fc1(x))
        x = self.fc2(x)
        return x  

# 创建网络并将网络部署到GPU上
net = Net().to(device)

上面的代码实现了对于训练网络的构建，主要是使用了两层的卷积网络进行特征的提取，然后将features导入到两层全连接linear网络层进行分类。

然后利用下面代码可以查看网络的参数数量构成。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
params = list(net.parameters())
print(len(params))
for i in range(len(params)):
    print(params[i].size())

# 8
# torch.Size([6, 1, 5, 5])
# torch.Size([6])
# torch.Size([16, 6, 5, 5])
# torch.Size([16])
# torch.Size([120, 1024])
# torch.Size([120])
# torch.Size([10, 120])
# torch.Size([10])


网络构建完成之后就需要对优化器进行设置，这里我选择使用交叉熵损失函数以及Adam优化方式。learning-rate设置为0.001。

1
2
3
import torch.optim as optim
criterion= nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(),lr=0.001)

网络和训练损失函数以及optimizer都设置好了，可以开始网络的训练过程了，我们的训练次数一共为20次。训练网络部分的代码如下所示：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

'''
    train the neural network
'''
for epoch in range(epoch_time):
    running_loss = 0.0
    for i, data in enumerate(trainloader):
        # 注意这里要将images和labels的数据放到GPU上。
        images,labels = data[0].to(device),data[1].to(device)

        # print(labels)
        optimizer.zero_grad()

        predicted = net(images)

        loss = criterion(predicted,labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i%100 == 99:
            print("[%d / %5d] loss: %3f" %(epoch+1,i+1,running_loss/100))
            running_loss = 0.0

print("finished training!")

torch.save(net.state_dict(),path)

训练完成之后将网络参数保存到预先设置好的文件中。

接下来进行网络的测试，我们是用with torch.no_grad():包裹测试代码，使得torch不跟踪运算过程，避免进一步计算梯度降低运算性能。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

net = Net().to(device)
net.load_state_dict(torch.load(path))

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for i ,data in enumerate(testloader):
        images ,labels = data[0].to(device),data[1].to(device)

        outputs = net(images)
        print(outputs,labels)

        _,predicted = torch.max(outputs,1)
        c = (predicted==labels).squeeze()
        loss = criterion(outputs,labels)
        print(loss)
        # print(c,predicted,labels)
        for i in range(len(c)):
            label = labels[i]
            class_correct[label]+= c[i].item()
            class_total[label]+=1
    
for i in range(10):
    print("accuracy of Number:%d : %.2f %%"%(i,100*class_correct[i]/class_total[i]))

训练结果如下：

1
2
3
4
5
6
7
8
9
10
accuracy of Number:0 : 99.39 %
accuracy of Number:1 : 99.74 %
accuracy of Number:2 : 99.22 %
accuracy of Number:3 : 99.21 %
accuracy of Number:4 : 99.08 %
accuracy of Number:5 : 99.10 %
accuracy of Number:6 : 98.75 %
accuracy of Number:7 : 99.03 %
accuracy of Number:8 : 98.56 %
accuracy of Number:9 : 98.81 %

可以看到对手写数字的识别性能还是很不错的，基本都能达到98%以上。并且这个网络还可以进一步提高。