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