pytorch神经网络搭建-手写字符识别

import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data
import matplotlib.pyplot as plt
from torch.autograd import Variable
 
# 超参数定义
EPOCH = 5  #定义太大的话CPU计算时间会太长 train the training data n times
LR = 0.01
BATCH_SIZE = 50

DOWNLOAD_MNIST = True  #没下载好的时候是True 下载好了改成 False
 
# 下载MNIST数据集
train_data = torchvision.datasets.MNIST(
    root = './mnist/',
    # 是否是训练数据
    train = True,
    # 数据变换(0, 255) -> (0, 1) 把下载的数据改成TENSOR
    transform = torchvision.transforms.ToTensor(),
    # 是否下载MNIST数据
    download = DOWNLOAD_MNIST
)
 
test_data = torchvision.datasets.MNIST(
    root = './mnist/',
    # 是否是训练数据
    train = False,
    # 数据变换(0, 255) -> (0, 1)
    transform = torchvision.transforms.ToTensor(),
    # 是否下载MNIST数据
    download = DOWNLOAD_MNIST
)
# 查看图像
plt.imshow(train_data.train_data[0].numpy(), cmap = 'gray')
plt.title('%i' % train_data.train_labels[0])
plt.show()

plt.imshow(test_data.test_data[0].numpy(), cmap = 'gray')
plt.title('%i' % test_data.test_labels[0])
plt.show()
# 数据加载
train_loader = Data.DataLoader(dataset = train_data, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)
#test_x = Variable(torch.unsqueeze(test_data.test_data,dim=1),volatile=True).type(torch.FloatTensor)[:2000]/255.
#test_y = test_data.test_lable[:2000] #测试只取了前两千个 为了节省时间 这里和莫烦视频里的不一样
test_loader = Data.DataLoader(dataset = test_data, batch_size = BATCH_SIZE, shuffle = False, num_workers = 1)
 

 
print(train_data.__len__)
print(test_data.__len__)
print(train_loader.__len__)
print(test_loader.__len__)
 
 
# 定义卷积神经网络 （重点）
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(      #卷积层
            nn.Conv2d(      #(1,28,28)     #卷积层相当于一个三维过滤器(有高度)(高度用来提取的特征属性)
                in_channels = 1,             #图片有多少个层 黑白图：1层  彩图:RGB 3层
                out_channels = 16,           #16个过滤器 同时提取16个特征
                kernel_size = 5,             #过滤器的大小 5*5
                stride = 1,                  #每隔一步跳一下
                padding = 2                  #扫到边缘过滤器超出范围时，多加一圈为0的信息
                                             #padding 的值的计算  if stride=1,padding=(kernel_size-1)/2=(5-1)/2=2
            ),#-->(16,28,28)
            nn.ReLU(),  #-->(16,28,28)   #神经网络
            nn.MaxPool2d(kernel_size = 2)#池化层
               #-->(16,14,14)               #使用了一个2*2的过滤器，选择2*2区域中的最大值，相当于把原始图片裁剪了一下，
                                            #换成长宽变小，高度不变
        )
        # conv1输出为(16, 14, 14)
        self.conv2 = nn.Sequential(  
            nn.Conv2d(16, 32, 5, 1, 2), #-->(32,14,14)
            nn.ReLU(),                  #-->(32,14,14)
            nn.MaxPool2d(2)             #-->(32,7,7)
            #Pooling 有两种MaxPool2d,AvgPool2d(平均值) 一般选MAX的
        )
        # conv2输出为(32, 7, 7)
        self.output = nn.Linear(32 * 7 * 7, 10) 
 
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)   #(batch,32,7,7)  考虑了batch
        x = x.view(x.size(0), -1)   #(batch,32*7*7)
        prediction = self.output(x)  #这样就有了上面的参数值
        return prediction
 
cnn = CNN()
print (cnn)
# 定义优化器
optimizer = torch.optim.Adam(cnn.parameters(), lr = LR, betas= (0.9, 0.999))
 
# 定义损失函数
loss_func = nn.CrossEntropyLoss()
 
# 训练
for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):
       
        # print('Epoch: ', epoch, 'Step: ', step) 
        x_var = Variable(x)
        y_var = Variable(y)
        
        prediction = cnn(x_var)
        loss = loss_func(prediction, y_var)
       # print(loss)

        optimizer.zero_grad() # 清空上一步的残余更新参数值
        loss.backward()     # 误差反向传播, 计算参数更新值
        optimizer.step()    # 将参数更新值施加到 net 的 parameters 上
 
        if step % 50 == 0:    
            correct = 0.0
            for step_test, (test_x, test_y) in enumerate(test_loader):
                # print(step_test)

                test_x = Variable(test_x)
                test_output = cnn(test_x)

                pred_y = torch.max(test_output, 1)[1].data.squeeze()
                correct += sum(pred_y == test_y)

            accuracy = correct / test_data.test_data.size(0)
            print ('epoch: ', epoch, '| train loss: %.4f' % loss.data, '| accuracy: ', accuracy)