import torch
from matplotlib import pyplot as plt

# 1.准备数据
# y = 3x +0.8
learning_rate = 0.01
x = torch.randn([500,1])
y_true = 3*x+0.8

# 2.计算预测值
w = torch.rand([1,1],requires_grad=True)
b = torch.tensor(0,requires_grad=True,dtype=torch.float32)

# 4.通过循环，反向传播，更新参数
for i in range(2000):
    # 3.计算损失
    y_predict = torch.matmul(x, w) + b
    loss = (y_predict - y_true).pow(2).mean()


    # 每次反向传播前把梯度置为0
    if w.grad is not None:
        w.grad.data.zero_()
    if b.grad is not None:
        b.grad.data.zero_()

    loss.backward() # 反向传播

    w.data -= learning_rate * w.grad.data
    b.data -= learning_rate * b.grad.data
    print("w,b,loss",w.item(),b.item(),loss.item())

plt.figure(figsize=(20,8))
plt.scatter(x.numpy().reshape(-1),y_true.numpy().reshape(-1))
y_predict = torch.matmul(x,w)+b
plt.plot(x.numpy().reshape(-1),y_predict.detach().numpy().reshape(-1),c="r")
plt.show()