測試結果
最後兩行分別爲預測類別與真實類別。
數據預覽
這裏的數據使用的是mnist數據集,大家可以將代碼中的DOWNLOAD_MNIST值修改爲True進行自動下載。
代碼
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")