pytorch實現MNIST手寫體識別

阿新 • • 發佈：2020-02-15

本文例項為大家分享了pytorch實現MNIST手寫體識別的具體程式碼，供大家參考，具體內容如下

實驗環境

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,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,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部分截圖如下：

pytorch實現MNIST手寫體識別

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%

pytorch實現MNIST手寫體識別