神经网络拟合函数

1.准备数据

导入需要使用的第三方库

import torch
import numpy as np
import random
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.functional as func

根据给定的w、b生成训练集和测试集(按9:1划分)

def gen_data(w, b, num):
    x = torch.linspace(0, 10, num, dtype=torch.float32).reshape(-1, 1)
    print(x.shape)
    y = torch.matmul(x, w) + b
    y += torch.normal(0, 0.01, y.shape)
    print(y.shape)
    
    train_x = x[: num * 9 // 10]
    train_y = y[: num * 9 // 10]
    
    test_x = x[num * 9 // 10 :]
    test_y = y[num * 9 // 10 :]
    
    return train_x, train_y, test_x, test_y

注意这里为什么要reshape，主要原因是要进行矩阵乘法，需要将向量转化成二维张量。

train_x, train_y, test_x, test_y = gen_data(torch.tensor([1.0]).reshape(1,1), 1.0, 100)

2.建立模型

拟合一次函数

class model(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(1, 1)
        
    def forward(self, x):
        x = self.hidden(x)
        return x

实现forward方法即可

3.训练模型

输入参数很简单，不用解释了

def train(train_x, train_y, epochs, opt, loss_func, net):
    for epoch in range(epochs):
        pred = net(train_x)
        loss = loss_func(pred, train_y)
        
        opt.zero_grad()
        loss.backward()
        opt.step()
        
        if epoch % 500 == 0:
            print("loss: ", loss)
    return list(net.parameters())

4.设置超参数

loss_func = torch.nn.MSELoss()
opt = torch.optim.SGD(net.parameters(), lr = 0.01)
epochs = 10000

这里选择均方误差损失函数，优化器选择随机梯度下降法，训练轮次为10000

5.开始训练

net = model()
train(train_x, train_y, epochs, opt, loss_func, net)

6.预测结果

def prediction(test_x, test_y, net):
    pred = net(test_x)
    with torch.no_grad():
        loss = loss_func(pred, test_y)
    plt.plot(test_x.numpy(), test_y.numpy(), 'o')
    plt.plot(test_x.numpy(), pred.detach().numpy())
    plt.show()
    print("pred_loss: ", loss)

prediction(test_x, test_y, net)

上图可以看出，误差已经很小了

7.拟合三次多项式函数

生成数据集

def gen_data(a, b, c, d, num):
    x = torch.linspace(0, 1, num, dtype=torch.float32).reshape(-1, 1)
    y = torch.matmul(x ** 3, a) + torch.matmul(x ** 2, b) + torch.matmul(x, c) + d
    y += torch.normal(0, 0.01, y.shape)
    plt.plot(x.numpy(), y.numpy(), 'o')
    plt.show()
    
    train_x = x[: num * 9 // 10]
    train_y = y[: num * 9 // 10]
    
    test_x = x[num * 9 // 10 :]
    test_y = y[num * 9 // 10 :]
    
    return train_x, train_y, test_x, test_y

可以利用Sequential创建神经网络

net = nn.Sequential(nn.Linear(1,5),
                    nn.ReLU(),
                    nn.Linear(5, 10),
                    nn.ReLU(),
                    nn.Linear(10, 5),
                    nn.ReLU(),
                    nn.Linear(5, 1)
                   )

拟合结果如下，图一是对训练数据的拟合效果，图二是对测试数据的拟合效果