Pytorch模型训练笔记(二)`Mnist手写数字识别(源码贡献)+神经网络模型CNN卷积框架知识图

书接上回·还差一些内容就Pytorch第一个模型训练+CNN卷积项目跑通了.先看知识图

 

话不多说直接代码

import torch
import numpy as np
import os
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.nn as nn
from torch.optim import Adam
from torchvision.datasets import MNIST
from torchvision.transforms import Compose,ToTensor,Normalize
BATCH_SIZE = 128
TEST_BATCH_SIZE=1000
#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,transform=transform_fn,download=True)
    data_loader = DataLoader(dataset,batch_size=BATCH_SIZE,shuffle=True)
    return data_loader
#2.构建模型
class MnistMode(nn.Module):
    def __init__(self):
        super(MnistMode, self).__init__()
        self.fc1 = nn.Linear(1*28*28,28)
        self.fc2 = nn.Linear(28,10)



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

        return F.log_softmax(out)
model = MnistMode()
optimizer = Adam(model.parameters(),lr=0.001)
if os.path.exists("./model/model.pkl"):#如果路径存在
    #加载训练好的模型文件
    model.load_state_dict(torch.load("./model/model.pkl"))
    #加载训练好的模型文件
    optimizer.load_state_dict(torch.load("./model/optimizer.pkl"))
def train(epoch):
    """实现训练过程"""
    data_loader = get_dataloader()
    for idx,(input,traget) in enumerate(data_loader):
        optimizer.zero_grad()
        output=model(input)#调用模型得到预测值
        loss=F.nll_loss(output,traget)#得到损失
        loss.backward()#反向传播
        optimizer.step()#梯度更新
        if idx%10==0:
            print(epoch,idx,loss.item())
        #模型的保存
        if idx%100==0:
            torch.save(model.state_dict(),"./model/model.pkl")
            torch.save(optimizer.state_dict(), "./model/optimizer.pkl")
def test():#测试集
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(train=False,batch_size=TEST_BATCH_SIZE)
    for idx,(input,traget) in enumerate(test_dataloader):
        with torch.no_grad():
            output=model(input)
            cur_loss = F.nll_loss(output,traget)
            loss_list.append(cur_loss)
            #计算准确率
            pred=output.max(dim=-1)[-1]
            cur_acc=pred.eq(traget).float().mean()
            acc_list.append(cur_acc)
    print('平均准确,平均损失',np.mean(acc_list),np.mean(loss_list))


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

 在右侧选择或者自行添加Anaconda pytorch环境下 在Pycharm运行