56 lines
1.2 KiB
Python
56 lines
1.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from torch import optim
|
|
|
|
# x 张量
|
|
x = torch.linspace(-10, 10, 100).unsqueeze(1)
|
|
|
|
# 特征矩阵
|
|
X = torch.cat([x, x**2, x**3, torch.sin(2*x)], dim=1)
|
|
|
|
# 高斯噪声
|
|
noise = torch.from_numpy(np.random.normal(0, 3, size=x.shape)).float()
|
|
|
|
# 带噪声的 y
|
|
y_noisy = 0.5 * x**3 - 2 * x**2 + 3*x + 5 + 4 * torch.sin(2*x) + noise
|
|
|
|
# 可视化
|
|
'''
|
|
plt.scatter(x, y_noisy, label="Noisy data")
|
|
plt.plot(x, 0.5 * x**3 - 2 * x**2 + 3*x + 5 + 4 * torch.sin(2*x), color='red', label="Original function")
|
|
plt.legend()
|
|
plt.show()
|
|
'''
|
|
ez = nn.Sequential(
|
|
nn.Linear(4, 8),
|
|
nn.ReLU(),
|
|
nn.Linear(8, 4),
|
|
nn.ReLU(),
|
|
nn.Linear(4, 1),
|
|
)
|
|
|
|
criterion = nn.MSELoss()
|
|
optimizer = optim.Adam(ez.parameters(), lr=0.0065)
|
|
|
|
for i in range(10000):
|
|
#向前传播
|
|
out = ez(X)
|
|
loss = criterion(out, y_noisy)
|
|
#反向传播
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if i % 100 == 0:
|
|
print("Step:", i, "loss =", loss.item())
|
|
|
|
with torch.no_grad():
|
|
y_pred = ez(X)
|
|
|
|
plt.scatter(x.numpy(), y_noisy.numpy(), alpha=0.3, label="Noisy data")
|
|
plt.plot(x.numpy(), y_pred.numpy(), color='red', label="NN prediction")
|
|
plt.legend()
|
|
plt.show()
|