神经网络拟合二元函数_pytorch实现

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

class Neuro_net(torch.nn.Module):
    """搭建神经网络"""
    def __init__(self):
        super(Neuro_net, self).__init__()   # 继承__init__功能
        self.hidden_layer1 = torch.nn.Linear(2, 100)
        self.hidden_layer2 = torch.nn.Linear(100, 100)
        self.output_layer = torch.nn.Linear(100, 1)

    def forward(self, x):
        x = self.hidden_layer1(x)
        x = F.relu(x)
        c = x
        for i in range(3):
            x = self.hidden_layer2(x)
            x = F.relu(x)

        pridect_y = self.output_layer(x)
        return pridect_y

train_data = np.zeros((10000,3))
for i in range(10000):
    train_data[i][0] = random.uniform(-1, 1)
    train_data[i][1] = random.uniform(-1, 1)
    train_data[i][2] = train_data[i][0]**2 + train_data[i][1]**2
x_data = train_data[:, 0:2]
y_data = train_data[:, 2].reshape(10000, 1)
print(x_data.shape, y_data.shape)

net = Neuro_net()
# optimizer 优化
optimizer = torch.optim.SGD(net.parameters(), lr=0.2)
# loss funaction
loss_funaction = torch.nn.MSELoss()
epoch = 500
x_data = torch.tensor(x_data, dtype=torch.float32)
y_data = torch.tensor(y_data, dtype=torch.float32)


plt.ion()
for step in range(epoch):
    pridect_y = net(x_data)  # 喂入训练数据 得到预测的y值
    loss = loss_funaction(pridect_y, y_data)  # 计算损失

    optimizer.zero_grad()    # 为下一次训练清除上一步残余更新参数
    loss.backward()          # 误差反向传播，计算梯度
    optimizer.step()         # 将参数更新值施加到 net 的 parameters 上

    if step % 100 == 0:
        print("已训练{}步 | loss：{}.".format(step, loss))
        plt.cla()
        ax = plt.subplot(111, projection='3d')
        ax.scatter(x_data[:, 0], x_data[:, 1], y_data, c='g')
        ax.scatter(x_data[:, 0], x_data[:, 1], pridect_y.data.numpy(), c='r')
        plt.pause(0.1)


plt.ioff()
plt.show()

test_data = np.zeros((2000, 4))
for i in range(2000):
    test_data[i][0] = random.uniform(-1, 1)
    test_data[i][1] = random.uniform(-1, 1)
    test_data[i][2] = test_data[i][0]**2 + test_data[i][1]**2
    
x_test = test_data[:, 0:2]
x_test = torch.tensor(x_test, dtype=torch.float32)
y_test = net(x_data)

ax = plt.subplot(111, projection='3d')
ax.scatter(x_data[:, 0].data.numpy(), x_data[:, 1].data.numpy(), y_test[:, 0].data.numpy(), c='r')
plt.show()

程序运行过程

图片截的不好，凑合着看吧。