最後兩行分別爲預測類別與真實類別。數據庫
這裏的數據使用的是mnist數據集,你們能夠將代碼中的DOWNLOAD_MNIST值修改成True進行自動下載。ide
import torch import torch.nn as nn import torch.utils.data as Data import torchvision # 數據庫模塊 import matplotlib.pyplot as plt #訓練整批數據多少次,這裏爲了節約時間,只訓練一次 EPOCH=1 #每次批處理50個數據 BATCH_SIZE=50 #學習效率 LR=0.001 # 若是已經下載好了mnist數據就寫上False DOWNLOAD_MNIST = False #訓練的數據集:Mnist手寫數字 train_data=torchvision.datasets.MNIST( #保存或提取數據集的位置 root='./mnist/', #該數據是訓練數據 train=True, #轉換PIL.Image or numpy.ndarray成torch.FloatTensor (C x H x W), 訓練的時候 normalize 成 [0.0, 1.0] 區間 transform=torchvision.transforms.ToTensor(), #沒下載就下載,下載了就不用再下了 download=DOWNLOAD_MNIST, ) #繪製一下數據集 #黑色的地方的值都是0, 白色的地方值大於0. print(train_data.train_data.size()) # (60000, 28, 28) print(train_data.train_labels.size()) # (60000) plt.imshow(train_data.train_data[2].numpy(), cmap='gray') plt.title('%i' % train_data.train_labels[2]) plt.show() #測試數據 test_data=torchvision.datasets.MNIST(root='./mnist/',train=False) #批訓練50samples,1 channel,28x28 (50, 1, 28, 28) train_loader=Data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True) #這裏只測試了前2000個 #特徵 test_x=torch.unsqueeze(test_data.test_data,dim=1).type(torch.FloatTensor)[:2000]/255. #標籤 test_y=test_data.test_labels[:2000] #構建CNN模型 class CNN(nn.Module): def __init__(self): super(CNN,self).__init__() #input shape(1,28,28) self.conv1=nn.Sequential( #卷積 nn.Conv2d( in_channels=1, out_channels=16, #filter size kernel_size=5, #filter movement/step stride=1, #若是想要con2d出來的圖片長寬沒有變化, #padding=(kernel_size-1)/2當stride=1 padding=2, ), #output shape(16,28,28) #激勵函數 nn.ReLU(), #池化 # 在2x2空間裏向下採樣,output shape(16,14,14) nn.MaxPool2d(kernel_size=2), ) #input shape(16,14,14) self.conv2=nn.Sequential( nn.Conv2d(16,32,5,1,2), #output shape(32,14,14) #激勵函數 nn.ReLU(), #output shape(32,7,7) nn.MaxPool2d(2), ) #全鏈接層——進行分類。這裏將其分紅了10類 self.out=nn.Linear(32*7*7,10) def forward(self,x): x=self.conv1(x) x=self.conv2(x) #展平多維的卷積圖成(batch_size,32*7*7) x=x.view(x.size(0),-1) output=self.out(x) return output cnn=CNN() print(cnn) #訓練 #優化器 optimizer=torch.optim.Adam(cnn.parameters(),lr=LR) #損失函數 loss_func=nn.CrossEntropyLoss() #開始訓練 for epoch in range(EPOCH): for step,(b_x,b_y) in enumerate(train_loader): #將數據輸入nn而且獲得output output=cnn(b_x) #計算output與真實值之間的偏差 loss=loss_func(output,b_y) #清空上一步殘餘更新參數值 optimizer.zero_grad() #偏差反向傳播,讓參數進行更新 loss.backward() #將更新後的參數值施加到nn的parameters上 optimizer.step() #測試:選取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') # if __name__=='__main__': # print("hello word")