利用torch.nn实现前馈神经网络解决多分类问题

6、利用torch.nn实现前馈神经网络解决多分类问题

#导入必要的包
import torch
import numpy as np
from torch import nn
from torchvision.datasets import MNIST
import torchvision.transforms  as transforms
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torch import nn

#读取数据
mnist_train = datasets.MNIST(root = './data',train = True,download = False,transform =transforms.ToTensor())
mnist_test = datasets.MNIST(root ='./data',train = False,download = False,transform = transforms.ToTensor())

#训练集
batch_size = 256
train_iter = DataLoader( 
    dataset = mnist_train,
    shuffle = True,
    batch_size = batch_size,
    num_workers = 0
)
#测试集
test_iter = DataLoader(
    dataset  = mnist_test,
    shuffle  =False,
    batch_size = batch_size,
    num_workers = 0
)

#定义模型
num_input,num_hidden1,num_hidden2,num_output = 28*28,512,256,10

class DNN(nn.Module):
    def __init__(self,num_input,num_hidden1,num_hidden2,num_output):
        super(DNN,self).__init__()
        self.linear1 = nn.Linear(num_input,num_hidden1)
        self.linear2 = nn.Linear(num_hidden1,num_hidden2)
        self.linear3 = nn.Linear(num_hidden2,num_output)
    def forward(self,input):
        input = input.view(-1,784)
        out = self.linear1(input)
        out = self.linear2(out)
        out = self.linear3(out)
        return out

#参数初始化
net = DNN(num_input,num_hidden1,num_hidden2,num_output)
for param in net.parameters():
    nn.init.normal_(param,mean=0,std=0.001)

#定义训练函数
def train(net,train_iter,test_iter,loss,num_epochs):
    train_ls,test_ls,train_acc,test_acc = [],[],[],[]
    for epoch in range(num_epochs):
        train_ls_sum,train_acc_sum,n = 0,0,0
        for x,y in train_iter:
            y_pred = net(x)
            l = loss(y_pred,y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            train_ls_sum +=l.item()
            train_acc_sum += (y_pred.argmax(dim = 1)==y).sum().item()
            n += y_pred.shape[0]
        train_ls.append(train_ls_sum)
        train_acc.append(train_acc_sum/n)
        
        test_ls_sum,test_acc_sum,n = 0,0,0
        for x,y in test_iter:
            y_pred = net(x)
            l = loss(y_pred,y)
            test_ls_sum +=l.item()
            test_acc_sum += (y_pred.argmax(dim = 1)==y).sum().item()
            n += y_pred.shape[0]
        test_ls.append(test_ls_sum)
        test_acc.append(test_acc_sum/n)
        print('epoch %d, train_loss %.6f,test_loss %f, train_acc %.6f,test_acc %f'
              %(epoch+1, train_ls[epoch],test_ls[epoch], train_acc[epoch],test_acc[epoch]))
    return train_ls,test_ls,train_acc,test_acc

#训练次数和学习率
num_epochs = 20
lr = 0.1
loss  = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(),lr=lr)

#训练模型
train_loss,test_loss,train_acc,test_acc = train(net,train_iter,test_iter,loss,num_epochs)

#结果可视化
x = np.linspace(0,len(train_loss),len(train_loss))
plt.plot(x,train_loss,label="train_loss",linewidth=1.5)
plt.plot(x,test_loss,label="test_loss",linewidth=1.5)
plt.xlabel("epoch")
plt.ylabel("loss")
plt.legend()
plt.show()