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

  1 import torch
  2 import torchvision
  3 import torch.nn as nn
  4 import torch.utils.data as Data
  5 import matplotlib.pyplot as plt
  6 from torch.autograd import Variable
  7  
  8 # 超参数定义
  9 EPOCH = 5  #定义太大的话CPU计算时间会太长 train the training data n times
 10 LR = 0.01
 11 BATCH_SIZE = 50
 12 
 13 DOWNLOAD_MNIST = True  #没下载好的时候是True 下载好了改成 False
 14  
 15 # 下载MNIST数据集
 16 train_data = torchvision.datasets.MNIST(
 17     root = ‘./mnist/‘,
 18     # 是否是训练数据
 19     train = True,
 20     # 数据变换(0, 255) -> (0, 1) 把下载的数据改成TENSOR
 21     transform = torchvision.transforms.ToTensor(),
 22     # 是否下载MNIST数据
 23     download = DOWNLOAD_MNIST
 24 )
 25  
 26 test_data = torchvision.datasets.MNIST(
 27     root = ‘./mnist/‘,
 28     # 是否是训练数据
 29     train = False,
 30     # 数据变换(0, 255) -> (0, 1)
 31     transform = torchvision.transforms.ToTensor(),
 32     # 是否下载MNIST数据
 33     download = DOWNLOAD_MNIST
 34 )
 35 # 查看图像
 36 plt.imshow(train_data.train_data[0].numpy(), cmap = ‘gray‘)
 37 plt.title(‘%i‘ % train_data.train_labels[0])
 38 plt.show()
 39 
 40 plt.imshow(test_data.test_data[0].numpy(), cmap = ‘gray‘)
 41 plt.title(‘%i‘ % test_data.test_labels[0])
 42 plt.show()
 43 # 数据加载
 44 train_loader = Data.DataLoader(dataset = train_data, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)
 45 #test_x = Variable(torch.unsqueeze(test_data.test_data,dim=1),volatile=True).type(torch.FloatTensor)[:2000]/255.
 46 #test_y = test_data.test_lable[:2000] #测试只取了前两千个 为了节省时间 这里和莫烦视频里的不一样
 47 test_loader = Data.DataLoader(dataset = test_data, batch_size = BATCH_SIZE, shuffle = False, num_workers = 1)
 48  
 49 
 50  
 51 print(train_data.__len__)
 52 print(test_data.__len__)
 53 print(train_loader.__len__)
 54 print(test_loader.__len__)
 55  
 56  
 57 # 定义卷积神经网络 （重点）
 58 class CNN(nn.Module):
 59     def __init__(self):
 60         super(CNN, self).__init__()
 61         self.conv1 = nn.Sequential(      #卷积层
 62             nn.Conv2d(      #(1,28,28)     #卷积层相当于一个三维过滤器(有高度)(高度用来提取的特征属性)
 63                 in_channels = 1,             #图片有多少个层 黑白图：1层  彩图:RGB 3层
 64                 out_channels = 16,           #16个过滤器 同时提取16个特征
 65                 kernel_size = 5,             #过滤器的大小 5*5
 66                 stride = 1,                  #每隔一步跳一下
 67                 padding = 2                  #扫到边缘过滤器超出范围时，多加一圈为0的信息
 68                                              #padding 的值的计算  if stride=1,padding=(kernel_size-1)/2=(5-1)/2=2
 69             ),#-->(16,28,28)
 70             nn.ReLU(),  #-->(16,28,28)   #神经网络
 71             nn.MaxPool2d(kernel_size = 2)#池化层
 72                #-->(16,14,14)               #使用了一个2*2的过滤器，选择2*2区域中的最大值，相当于把原始图片裁剪了一下，
 73                                             #换成长宽变小，高度不变
 74         )
 75         # conv1输出为(16, 14, 14)
 76         self.conv2 = nn.Sequential(  
 77             nn.Conv2d(16, 32, 5, 1, 2), #-->(32,14,14)
 78             nn.ReLU(),                  #-->(32,14,14)
 79             nn.MaxPool2d(2)             #-->(32,7,7)
 80             #Pooling 有两种MaxPool2d,AvgPool2d(平均值) 一般选MAX的
 81         )
 82         # conv2输出为(32, 7, 7)
 83         self.output = nn.Linear(32 * 7 * 7, 10) 
 84  
 85     def forward(self, x):
 86         x = self.conv1(x)
 87         x = self.conv2(x)   #(batch,32,7,7)  考虑了batch
 88         x = x.view(x.size(0), -1)   #(batch,32*7*7)
 89         prediction = self.output(x)  #这样就有了上面的参数值
 90         return prediction
 91  
 92 cnn = CNN()
 93 print (cnn)
 94 # 定义优化器
 95 optimizer = torch.optim.Adam(cnn.parameters(), lr = LR, betas= (0.9, 0.999))
 96  
 97 # 定义损失函数
 98 loss_func = nn.CrossEntropyLoss()
 99  
100 # 训练
101 for epoch in range(EPOCH):
102     for step, (x, y) in enumerate(train_loader):
103        
104         # print(‘Epoch: ‘, epoch, ‘Step: ‘, step) 
105         x_var = Variable(x)
106         y_var = Variable(y)
107         
108         prediction = cnn(x_var)
109         loss = loss_func(prediction, y_var)
110        # print(loss)
111 
112         optimizer.zero_grad() # 清空上一步的残余更新参数值
113         loss.backward()     # 误差反向传播, 计算参数更新值
114         optimizer.step()    # 将参数更新值施加到 net 的 parameters 上
115  
116         if step % 50 == 0:    
117             correct = 0.0
118             for step_test, (test_x, test_y) in enumerate(test_loader):
119                 # print(step_test)
120 
121                 test_x = Variable(test_x)
122                 test_output = cnn(test_x)
123 
124                 pred_y = torch.max(test_output, 1)[1].data.squeeze()
125                 correct += sum(pred_y == test_y)
126 
127             accuracy = correct / test_data.test_data.size(0)
128             print (‘epoch: ‘, epoch, ‘| train loss: %.4f‘ % loss.data, ‘| accuracy: ‘, accuracy)
129 
130