pytorch数据集MNIST训练与测试实例

import os
import torch
import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose,ToTensor,Normalize
from torch.optim import Adam
import torch.nn as nn
import torch.nn.functional as F


BATCH_SIZE = 128
TEST_BATCH_SIZE = 516
#1、准备数据集
def get_dataloader(train=True,batch_size=BATCH_SIZE):
    transform_fn = Compose([ToTensor(),Normalize(mean=(0.1307,),std=(0.3081,))]) #mean和std的形状和通道数相同
    dataset = MNIST(root='./data',train=train,download=False,transform=transform_fn)
    data_loader = DataLoader(dataset,batch_size=batch_size,shuffle=True)
    return data_loader
    # for i in enumerate(data_loader):
    #     print(i)

#2.构建模型
class MnisModel(nn.Module):
    def __init__(self):
        super(MnisModel,self).__init__()
        self.fc1 = nn.Linear(1*28*28,28)
        self.fc2 = nn.Linear(28,10)

    def forward(self,input):
        """
        input:[batch_size,1,28,28]
        """
        #1.修改形状
        x = input.view([input.size(0),1*28*28]) # input.view(-1,1*28*28)
        #2.进行全连接的操作
        x = self.fc1(x)
        #3.进行激活函数处理，形状不会发生变化
        x = F.relu(x)
        #4.输出层
        out = self.fc2(x)
        return F.log_softmax(out,dim=-1)

# 1.实例化模型
model = MnisModel()
#2.实例优化器类
optimizer = Adam(model.parameters(),lr=0.001)
if os.path.exists("./model/model.pt"):
    model.load_state_dict(torch.load("./model/model.pt"))  #加载模型
    optimizer.load_state_dict(torch.load("./model/optimizer.pt"))  #加载优化器

def train(epoch):
    """
    实现训练过程
    """
    #3.加载数据集，遍历
    data_loader = get_dataloader()
    for idex,(input,target) in enumerate(data_loader):
        optimizer.zero_grad()  #4.梯度置为0
        output = model(input)  #5.调用模型，得到预测值
        loss = F.nll_loss(output,target)  #6.计算损失
        loss.backward()  #7.反向传播
        optimizer.step()  #8.梯度的更新
        if idex % 100 == 0:
            print(loss.item())

        if idex % 100 == 0:
            torch.save(model.state_dict(),"./model/model.pt")  #保存模型参数
            torch.save(optimizer.state_dict(),"./model/optimizer.pt") #保存优化器参数

def test():  #测试数据
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(False,batch_size=TEST_BATCH_SIZE)  #获取测试数据集
    for idx,(input,target) in enumerate(test_dataloader):
        # print(idx,target,input)
        # break
        with torch.no_grad():
            output = model(input)
            cur_loss = F.nll_loss(output,target)
            loss_list.append(cur_loss)
            #计算准备率
            #output [batch_size,10] target:[batch_size]
            pred = output.max(dim = -1)[-1]
            cur_acc = pred.eq(target).float().mean()
            acc_list.append(cur_acc)

    print("平均准确率，平均损失",np.mean(acc_list),np.mean(loss_list))


if __name__ == '__main__':
    # for i  in range(3):  #训练三轮
    #     train(i)
    test()