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
