[pytorch入门]使用全连接神经网络实现MNIST手写数字识别

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')