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函数里面的注释取消