import torch
from torch import nn
from torch.utils.data import DataLoader
import numpy as np
from matplotlib import pyplot as plt
from load_mnist import load_mnist
#pytorch 中计计算交叉熵损失函数时,
#输入的 label 不能是 one-hot 格式,因为函数内部会自己处理成 one hot 格式。
#需要在label后面加上.long(),不然会报格式错误
#超参数
lr = 1e-4
batch_size = 32
epochs = 10
def adjust_lr(optim,lr):
for params in optim.param_groups:
params['lr'] = lr
class Net(nn.Module):
def __init__(self,input_dim,output_dim):
super().__init__()
self.activate_f = nn.ReLU(True)
self.input_dim = input_dim
self.output_dim = output_dim
self.layer1 = nn.Sequential(
nn.Linear(self.input_dim,64),
nn.BatchNorm1d(64),
self.activate_f)
self.layer2 = nn.Sequential(
nn.Linear(64,128),
nn.BatchNorm1d(128),
self.activate_f)
self.layer3 = nn.Sequential(
nn.Linear(128,self.output_dim),
nn.BatchNorm1d(self.output_dim),
self.activate_f)
def forward(self,x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
return x
if __name__ =='__main__':
#载入数据
(x_train,y_train),(x_test,y_test) = load_mnist(flatten=True,one_hot_label=False)
x_train = torch.autograd.Variable(torch.Tensor(x_train))#.cuda()
y_train = torch.autograd.Variable(torch.Tensor(y_train))#.cuda()
with torch.no_grad():#测试集不需要计算梯度,节约内存空间
x_test = torch.autograd.Variable(torch.Tensor(x_test))#.cuda()
y_test = torch.autograd.Variable(torch.Tensor(y_test))#.cuda()
#将数据放入数据库
train_data = torch.utils.data.TensorDataset(x_train,y_train)
test_data = torch.utils.data.TensorDataset(x_test,y_test)
#生成数据迭代器
train_loader = DataLoader(train_data,batch_size=batch_size,shuffle=True)
test_loader = DataLoader(test_data,batch_size=batch_size,shuffle=True)
#实例化模型
model = Net(28*28,10)#.cuda()
#定义损失函数和优化函数
cost = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters(),lr=lr,weight_decay=1e-5)
#定义精度的计算所需参数
total_steps = x_train.shape[0]/batch_size
accs = 0
#训练模型
for epoch in range(epochs):
if epoch == epochs*4/5:
lr /= 10
adjust_lr(optimizer, lr)
if epoch == epochs*9/10:
lr /= 10
adjust_lr(optimizer, lr)
for step,(batch_x,batch_y) in enumerate(train_loader):
#前向传播
output = model(batch_x)
loss = cost(output,batch_y.long())
accs += ((output.argmax(axis=1)==batch_y).sum()).tolist()/batch_size
#反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch[{}/{}],Loss = {:.4f},Accuracy = {:.4f}\n'.format(epoch+1,epochs,loss.data.tolist(),accs/total_steps))
#精度清零,重新计算
accs = 0
#评估模型
model.eval()
for step,(batch_x,batch_y) in enumerate(test_loader):
output = model(batch_x)
accs += ((output.argmax(axis=1)==batch_y).sum()).tolist()/batch_size
print('Test Accuracy:{:.4f}'.format(accs/(x_test.shape[0]/batch_size)))
#保存模型
torch.save(model,r'my_model')
[pytorch入门]使用全连接神经网络实现MNIST手写数字识别
最新推荐文章于 2023-06-05 19:08:40 发布