import torch
import torch.nn as nn
from torch.optim import SGD
import matplotlib.pyplot as plt
#0.准备数据
x = torch.rand([500,1])
y_true = 3*x + 0.8
#1.定义模型
class MyLinear(nn.Module):
    def __init__(self):
        super(MyLinear, self).__init__()
        self.linear = nn.Linear(1,1)
    def forward(self,x):
        out = self.linear(x)
        return out

#2.实例化模型，优化器实例化，loss实例化
my_linear = MyLinear()
optimizer = SGD(my_linear.parameters(),0.001)
loss_fn = nn.MSELoss()
#3.循环，进行梯度下降，参数的更新
for i in range(30000):
    y_predict = my_linear(x)
    loss = loss_fn(y_predict,y_true)
    optimizer.zero_grad()
    loss.backward()
    #参数更新
    optimizer.step()
    if i%100 == 0:
        params = list(my_linear.parameters())
        print(loss.item(),params[0].item(),params[1].item())
my_linear.eval()
predict = my_linear(x)
predict = predict.data.numpy()
plt.scatter(x.data.numpy(),y_true.data.numpy(),c='r')
plt.plot(x.data.numpy(),predict)
plt.show()