用pytorch做手寫數字識別，識別l率達97.8%

pytorch做手寫數字識別

效果如下：

 

工程目錄如下

 

第一步  數據獲取

下載MNIST庫，這個庫在網上，執行下面代碼自動下載到當前data文件夾下

from torchvision.datasets import MNIST
import torchvision

mnist = MNIST(root='./data',train=True,download=True)

print(mnist)
print(mnist[0])
print(len(mnist))
img = mnist[0][0]
img.show()

　　

dataset.py文件，讀取數據並做預處理

 

'''
准備數據集
'''

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import torchvision


def mnist_dataset(train):

    func = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=(0.1307,),std=(0.3081,))
    ])

    #1.准備Mnist數據集
    return MNIST(root='./data',train=train,download=False,transform=func)

def get_dataloader(train = True):
    mnist = mnist_dataset(train)
    return DataLoader(mnist,batch_size=128,shuffle=True)

if __name__ == '__main__':
    for (images,labels) in get_dataloader():
        print(images.size())
        print(labels.size())
        break

 

　　

 

models.py文件，定義訓練的模型類

'''
定義模型
'''

import torch.nn as  nn
import torch.nn.functional as F

class MnistModel(nn.Module):

    def __init__(self):
        super(MnistModel,self).__init__()
        self.fc1 = nn.Linear(1*28*28,100)
        self.fc2 = nn.Linear(100,10)

    def forward(self,image):
        image_viewd = image.view(-1,1*28*28) #[batch_size,1*28*28]
        fc1_out = self.fc1(image_viewd) #[batch_size,100]
        fc1_out_relu = F.relu(fc1_out) #[batch_size,100]
        out = self.fc2(fc1_out_relu) #[batch_size,10]

        return F.log_softmax(out,dim=-1)  #帶權損失計算交叉熵

 

cong.py文件，定義一些常亮，設置使用cpu還是GPU　　

'''
項目配置
'''

import torch

train_batch_size = 128
test_batch_size = 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

　　

train.py文件，模型訓練文件，保存模型

"""
進行模型的訓練
"""
from dataset import get_dataloader
from models import MnistModel
from torch import optim
import torch.nn.functional as F
import conf
from tqdm import tqdm
import numpy as np
import torch
import os
from test import eval

#1. 實例化模型，優化器，損失函數
model = MnistModel().to(conf.device)
optimizer = optim.Adam(model.parameters(),lr=1e-3)

#2. 進行循環，進行訓練
def train(epoch):
    train_dataloader = get_dataloader(train=True)
    bar = tqdm(enumerate(train_dataloader),total=len(train_dataloader))
    total_loss = []
    for idx,(input,target) in bar:
        input = input.to(conf.device)
        target = target.to(conf.device)
        #梯度置為0
        optimizer.zero_grad()
        #計算得到預測值
        output = model(input)
        #得到損失
        loss = F.nll_loss(output,target)
        #反向傳播，計算損失
        loss.backward()
        total_loss.append(loss.item())
        #參數的更新
        optimizer.step()
        #打印數據
        if idx%10 ==0 :
            bar.set_description_str("epcoh:{} idx:{},loss:{:.6f}".format(epoch,idx,np.mean(total_loss)))
            torch.save(model.state_dict(),"./models/model.pkl")
            torch.save(optimizer.state_dict(),"./models/optimizer.pkl")

if __name__ == '__main__':
    for i in range(10):
        train(i)
        eval()

 

test.py文件，模型測試文件，測試模型准確率　　

'''
進行模型評估
'''

from dataset import get_dataloader
from models import MnistModel
from torch import optim
import torch.nn.functional as F
import conf
from tqdm import tqdm
import numpy as np
import torch
import os

def eval():
    #實例化模型，優化器，損失函數
    model = MnistModel().to(conf.device)

    if os.path.exists("./models/model.pkl"):
        model.load_state_dict(torch.load("./models/model.pkl"))

    test_dataloader = get_dataloader(train=False)
    total_loss = []
    total_acc = []
    with torch.no_grad():
        for input, target in test_dataloader:  # 2. 進行循環，進行訓練
            input = input.to(conf.device)
            target = target.to(conf.device)
            # 計算得到預測值
            output = model(input)
            # 得到損失
            loss = F.nll_loss(output, target)
            # 反向傳播，計算損失
            total_loss.append(loss.item())

            # 計算准確率
            ###計算預測值
            pred = output.max(dim=-1)[-1]
            total_acc.append(pred.eq(target).float().mean().item())
    print("test loss:{},test acc:{}".format(np.mean(total_loss), np.mean(total_acc)))

# if __name__ == '__main__':
#     # for i in range(10):
#     #     train(i)
#     eval()