手写数字识别

python 复制代码
# 使用pytorch完成手写数字识别
# 准备dataloader
import os
import torch
from torch.utils.data import DataLoader
# 这里使用torch自带的MNIST数据集
from torchvision.datasets import MNIST
# 这里是为了将图片进行处理,调用这三个函数
from torchvision.transforms import Compose, ToTensor, Normalize
from torch import nn
import torch.nn.functional as F
from torch.optim import Adam
import numpy as np

BATCH_SIZE = 128
TEST_BATCH_SIZE = 1000


# 1. 准备数据集
def get_dataloader(train=True,batch_size=BATCH_SIZE):
    transform_fn = Compose([
        ToTensor(),     # 因为图片本身不是一个tensor,所以要to一下
        Normalize(mean=(0.1307,),std=(0.3081,))        # mean和std的形状和通道数相同,里面的数据mean和std都是给定的,直接搜

    ])
    dataset = MNIST(root="data",train=train,transform=transform_fn)
    data_loader = DataLoader(dataset,batch_size=BATCH_SIZE,shuffle=True)
    return data_loader

# 2. 构建模型
class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()
        self.fc1 = nn.Linear(1*28*28,28) #后面这个28表示经过变换后的形状为28而不是1*28*28了
        self.fc2 = nn.Linear(28,10) #最后只有10个数所以为10

    def forward(self,input):

        # 1.修改形状
        x = input.view(-1,28*28*1)   #对数据形状变形,-1表示该位置根据后面的形状自动调整
        # 2.进行全连接操作
        x = self.fc1(x)     #[batch_size,28]
        # 3. 进行激活函数处理,形状不发生变化
        x = F.relu(x)   #[batch_size,28]
        # 4. 输出层
        out = self.fc2(x)     #[batch_size,10]

        return F.log_softmax(out,dim=-1)


# 实例化模型
model = MnistModel()
# 优化器类
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"))


# 3. 实现训练的过程
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,target) in enumerate(test_dataloader):
        # 测试不需要梯度
        with torch.no_grad():
            output = model(input)
            cur_loss = F.nll_loss(output,target)
            loss_list.append(cur_loss)
            # 计算准确率
            # -1代表每一行的最大值,0代表每一列的最大值。里面有两组数,一组是值,一组是那个值的坐标.我们这里要坐标
            pred = output.max(dim=-1)[-1]
            # 比较两组数,eq方法返回的是bool值所以要变为float,再求均值
            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): # 训练3轮
    #     train(i)
    #
    test()

test和train是两个方向,想要走test就是现在代码写的,想运行train就是把main函数里面的注释取消

相关推荐
love_and_hope5 小时前
Pytorch学习--神经网络--完整的模型训练套路
人工智能·pytorch·python·深度学习·神经网络·学习
pzx_0016 小时前
【Pytorch】model.eval()与model.train()
人工智能·pytorch·机器学习
明晚十点睡7 小时前
2024AAAI | DiffRAW: 利用扩散模型从手机RAW图生成单反相机质量的RGB图像
人工智能·pytorch·python·深度学习
奔跑的小狗8 小时前
torch.nn.**和torch.nn.functional.**的区别
pytorch·深度学习·机器学习
MinIO官方账号1 天前
使用亚马逊 S3 连接器为 PyTorch 和 MinIO 创建地图式数据集
人工智能·pytorch·python
槿花Hibiscus1 天前
深度学习基础练习:从pytorch API出发复现LSTM与LSTMP
pytorch·深度学习·lstm
肥猪猪爸1 天前
VAE的原理及MNIST数据生成
人工智能·pytorch·深度学习·机器学习·计算机视觉·ai作画·vae
love_and_hope2 天前
Pytorch学习--神经网络--网络模型的保存与读取
人工智能·pytorch·python·深度学习·神经网络·学习
Shuai@2 天前
segmentation_pytorch库需要用huggingface,但是国内无法访问时,使用离线的方式
pytorch·python
小袁在上班2 天前
PyTorch 2.0: 新特性与升级指南
人工智能·pytorch·python