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