用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()