import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

# 生成数据
np.random.seed(0)
x = np.random.rand(100, 1) * 10  # 生成100个在0到10之间的随机数
y = 3.5 * x + 2 + np.random.randn(100, 1) * 2  # 真实的线性关系加上一些噪声

# 将numpy数组转换为torch张量
x_tensor = torch.from_numpy(x).float()
y_tensor = torch.from_numpy(y).float()


model = nn.Linear(in_features=1, out_features=1, bias=True)

criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

epochs = 100
for epoch in range(epochs):
    # 前向传播
    outputs = model(x_tensor)
    loss = criterion(outputs, y_tensor)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')

# 可视化原始数据和拟合结果
predicted = model(x_tensor).detach().numpy()  # detach()用于从计算图中移除张量，避免未来的梯度计算干扰。
plt.scatter(x, y, label='Original data')
plt.plot(x, predicted, color='red', label='Fitted line')
plt.legend()
plt.show()