利用torch.nn实现前馈神经网络解决回归问题

4、利用torch.nn实现前馈神经网络解决回归问题

#导入必要的包
import torch 
from torch import nn
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader,TensorDataset
from sklearn.model_selection import train_test_split
from collections import OrderedDict
from torch.nn import init

#创建数据集
num_input ,num_example = 500,10000
true_w = torch.ones(1,num_input)*0.0056
true_b = 0.028
x_data = torch.tensor(np.random.normal(0,0.001,size  = (num_example,num_input)),dtype = torch.float32)
y = torch.mm(x_data,true_w.t()) +true_b
y += torch.normal(0,0.001,y.shape)
train_x,test_x,train_y,test_y = train_test_split(x_data,y,shuffle= True,test_size=0.3)

#读取数据
batch_size = 50
train_dataset = TensorDataset(train_x,train_y)
train_iter = DataLoader(
    dataset = train_dataset,
    batch_size = batch_size,
    shuffle = True,
    num_workers = 0,
)
test_dataset = TensorDataset(test_x,test_y)
test_iter = DataLoader(
    dataset = test_dataset,
    batch_size = batch_size,
    shuffle = True,
    num_workers = 0,
)

#定义模型
model= nn.Sequential(OrderedDict([
    ('linear1',nn.Linear(num_input,256)),
    ('linear2',nn.Linear(256,128)),
    ('linear3',nn.Linear(128,1)),
])
)
#参数初始化
for param in model.parameters():
    init.normal_(param,mean = 0 ,std = 0.001)

#for param in model.state_dict():  #查看初始化参数
#    print(param)
#    print(model.state_dict()[param])

#定义学习率，损失函数，优化器
lr = 0.001
loss = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),lr)

#定义训练函数
def train(model,train_iter,test_iter,loss,num_epochs,batch_size,lr):
    train_ls,test_ls = [],[]
    for epoch in range(num_epochs):
        train_ls_sum ,test_ls_sum = 0,0
        for x,y in train_iter:
            y_pred = model(x)
            l = loss(y_pred,y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            train_ls_sum += l.item()
        for x ,y in test_iter:
            y_pred = model(x)
            l = loss(y_pred,y)
            test_ls_sum +=l.item()
        train_ls.append(train_ls_sum)
        test_ls.append(test_ls_sum)
        print('epoch %d,train_loss %.6f,test_loss %f'%(epoch+1, train_ls[epoch],test_ls[epoch]))
    return train_ls,test_ls

#开始训练
num_epochs = 30
train_loss ,test_loss = train(model,train_iter,test_iter,loss,num_epochs,batch_size,lr)

#结果可视化
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()