PyTorch学习CNN手写体识别

Posted 2021-06-07 My heart will go ~~

本节将介绍利用CNN进行手写体识别
首先呢，我们需要下载数据来进行训练。下载的代码如下：
注意：数据集下载一次就好，

DOWNLOAD_MNIST = True

train_data=torchvision.datasets.MNIST(#下载数据的代码
    root='./mnist',
    train=True,
    transform=torchvision.transforms.ToTensor(),  #(网上数据改为tensor),0-1之间，并复制到train_data中
    download=DOWNLOAD_MNIST#没有下载就=true，下载了就用false
    )

然后我们简单介绍一下卷及神经网络CNN的结构：

整体流程是：卷积(Conv2d) -> 激励函数(ReLU) -> 池化, 向下采样 (MaxPooling) -> 再卷积(Conv2d) -> 再激励函数(ReLU) -> 再池化, 向下采样 (MaxPooling) -> 展平多维的卷积成的特征图 -> 接入全连接层 (Linear) -> 输出。
全部代码如下：（给出了大部分的详细注释）

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

EPOCH=1  #训练整批数据的次数
BATCH_SIZE = 50#批训练的数据个数
LR = 0.001     # 学习率
DOWNLOAD_MNIST = False  


train_data=torchvision.datasets.MNIST(#下载数据的代码
    root='./mnist',
    train=True,
    transform=torchvision.transforms.ToTensor(),  #(网上数据改为tensor),0-1之间，并复制到train_data中
    download=DOWNLOAD_MNIST#没有下载就=true，下载了就用false
    )
    
    #plt测试一下下载的照片
# print(train_data.train_data.size())
# print(train_data.train_labels.size())
# plt.imshow(train_data.train_data[0].numpy(),cmap='gray')#呈现出第一张图片
# plt.title('%i'%train_data.train_labels[0])
# plt.show()

#批训练（50，1,28*28）
train_loader=Data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True)

#测试集,train说明不是traindata，是testdata
test_data=torchvision.datasets.MNIST(root='./mnist/',train=False)

# 为了节约时间, 我们测试时只测试前2000个
test_x = Variable(torch.unsqueeze(test_data.test_data, dim=1),volatile=True).type(torch.FloatTensor)[:2000]/255.   
test_y = test_data.test_labels[:2000]


#建立CNN网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        self.conv1=nn.Sequential(#卷积层1，包括以下三个内容
            nn.Conv2d(#卷基层，过滤器。
                in_channels=1,#图片的高度、层数。（因为输入的是二维图片，所以高度是1）
                out_channels=16,#16个filter的个数，同时进行扫描。输出的高度
                kernel_size=5,#kernel的宽、高都是5，5*5的扫描区域
                stride=1,#步长
                padding=2,#像素旁边一圈加上0的数据
                #如果 stride=1,padding=（kernel_size-1）/2=(5-1)/2=2
                ),
            nn.ReLU(),#激活函数,加了一层卷积层
            nn.MaxPool2d(kernel_size=2),#池化层，筛选重要的部分
            )
        self.conv2=nn.Sequential(#卷积层2
            nn.Conv2d(16,32,5,1,2),
            nn.ReLU(),#激活函数,加了一层卷积层
            nn.MaxPool2d(2),#池化层，筛选重要的部分
        )
        self.out=nn.Linear(32*7*7,10)#
        
    def forward(self,x):#展平的过程
        x=self.conv1(x)
        x=self.conv2(x)            #(batch,32,7,7)
        x=x.view(x.size(0),-1)     #(batch,32*7*7)
        output=self.out(x)
        return output

cnn=CNN()

optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)   # optimize all cnn parameters
loss_func = nn.CrossEntropyLoss() 

#训练过程
for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):   # 分配 batch data, normalize x when iterate train_loader
        b_x=Variable(x)
        b_y=Variable(y)

        output = cnn(b_x)               # cnn output
        loss = loss_func(output, b_y)   # cross entropy loss
        optimizer.zero_grad()           # clear gradients for this training step
        loss.backward()                 # backpropagation, compute gradients
        optimizer.step()                # apply gradients


        if step%50==0:
            test_output=cnn(test_x)
            pred_y=torch.max(test_output,1)[1].data.squeeze()
            accuracy=sum(pred_y==test_y)/test_y.size(0)

    
        

#最后取10组数据检查一下预测值到底对不对
test_output = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].numpy(), 'real number')