實驗環境
pytorch 1.4
Windows 10
python 3.7
cuda 10.1(我筆記本上沒有可以使用cuda的顯卡)
實驗過程
1. 確定我們要加載的庫
import torch
import torch.nn as nn
import torchvision #這裏面直接加載MNIST數據的方法
import torchvision.transforms as transforms # 將數據轉爲Tensor
import torch.optim as optim
import torch.utils.data.dataloader as dataloader
2. 加載數據
這裏使用所有數據進行訓練,再使用所有數據進行測試
train_set = torchvision.datasets.MNIST(
root='./data', # 文件存儲位置
train=True,
transform=transforms.ToTensor(),
download=True
)
train_dataloader = dataloader.DataLoader(dataset=train_set,shuffle=False,batch_size=100)# dataset可以省
'''
dataloader返回(images,labels)
其中,
images維度:[batch_size,1,28,28]
labels:[batch_size],即圖片對應的
'''
test_set = torchvision.datasets.MNIST(
root='./data',
train=False,
transform=transforms.ToTensor(),
download=True
)
test_dataloader = dataloader.DataLoader(test_set,batch_size=100,shuffle=False) # dataset可以省
3. 定義神經網絡模型
這裏使用全神經網絡作爲模型
class NeuralNet(nn.Module):
def __init__(self,in_num,h_num,out_num):
super(NeuralNet,self).__init__()
self.ln1 = nn.Linear(in_num,h_num)
self.ln2 = nn.Linear(h_num,out_num)
self.relu = nn.ReLU()
def forward(self,x):
return self.ln2(self.relu(self.ln1(x)))
4. 模型訓練
in_num = 784 # 輸入維度
h_num = 500 # 隱藏層維度
out_num = 10 # 輸出維度
epochs = 30 # 迭代次數
learning_rate = 0.001
USE_CUDA = torch.cuda.is_available() # 定義是否可以使用cuda
model = NeuralNet(in_num,h_num,out_num) # 初始化模型
optimizer = optim.Adam(model.parameters(),lr=learning_rate) # 使用Adam
loss_fn = nn.CrossEntropyLoss() # 損失函數
for e in range(epochs):
for i,data in enumerate(train_dataloader):
(images,labels) = data
images = images.reshape(-1,28*28) # [batch_size,784]
if USE_CUDA:
images = images.cuda() # 使用cuda
labels = labels.cuda() # 使用cuda
y_pred = model(images) # 預測
loss = loss_fn(y_pred,labels) # 計算損失
optimizer.zero_grad()
loss.backward()
optimizer.step()
n = e * i +1
if n % 100 == 0:
print(n,'loss:',loss.item())
訓練模型的loss部分截圖如下:
5. 測試模型
with torch.no_grad():
total = 0
correct = 0
for (images,labels) in test_dataloader:
images = images.reshape(-1,28*28)
if USE_CUDA:
images = images.cuda()
labels = labels.cuda()
result = model(images)
prediction = torch.max(result, 1)[1] # 這裏需要有[1],因爲它返回了概率還有標籤
total += labels.size(0)
correct += (prediction == labels).sum().item()
print("The accuracy of total {} images: {}%".format(total, 100 * correct/total))
實驗結果
最終實驗的正確率達到:98.22%