coding=utf-8
“”"
author:lei
function: 使用pytorch完成手写数字的识别
“”"
import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
import torch
import os
BATCH_SIZE = 512
TEST_BATCH_SIZE = 512
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=True, 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(12828, 28)
self.fc2 = nn.Linear(28, 10)
def forward(self, input):
"""
:param input: [batch_size, 1, 28, 28]
:return:
"""
# 1、修改形状 最好别直接写batch_size,如果batch_size进行改变,这里会出错;可用input.size[0]或者 -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 = 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"))
def train(epoch):
“”"
实现训练的过程
:param epoch:
:return:
“”"
data_loader = get_dataloader()
for idx, (input, target) in enumerate(data_loader):
output = model(input) # 调用模型,得到预测值
loss = F.nll_loss(output, target) # 得到损失
optimizer.zero_grad() # 梯度置为0
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():
test_dataloader = get_dataloader(train=False, batch_size=TEST_BATCH_SIZE)
loss_list = []
acc_list = []
for idx, (input, target) in enumerate(test_dataloader):
with torch.no_grad():
output = model(input) # 输出[batch_size, 10]
cur_loss = F.nll_loss(output, target)
# output [batch_size, 10] target:[batch_size]
pred = output.max(dim=-1)[-1] # 获取到最大值的位置
cur_acc = pred.eq(target).float().mean() # 将预测值和真实值进行比较,返回的结果为bool类型,将bool类型转换为浮点类型并求均值
acc_list.append(cur_acc) # 将每个batch放到acc_list中
loss_list.append(loss_list)
print("平均准确率,平均损失:", np.mean(acc_list), np.mean(loss_list)) # 求均值
if name == ‘main’:
# for i in range(3): # 训练三轮
# train(i)
# 进行测试
test()
# loader = get_dataloader(train=False)
#
# for input, label in loader:
# print(label.size()) # torch.Size([128])
# break