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

# 数据生成
x = np.linspace(-3, 3, 100).reshape(-1, 1)
y = 2 * x + 1 + np.random.normal(0, 0.5, x.shape)
x_tensor = torch.FloatTensor(x)
y_tensor = torch.FloatTensor(y)

# 模型定义
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    def forward(self, x):
        return self.linear(x)

# 训练配置
model = LinearModel()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练循环
for epoch in range(1000):
    pred = model(x_tensor)
    loss = criterion(pred, y_tensor)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# 结果输出
w = model.linear.weight.item()
b = model.linear.bias.item()
print(f'Final equation: y = {w:.2f}x + {b:.2f}')

# 可视化
plt.scatter(x, y)
plt.plot(x, w*x + b, 'r-')
plt.show()