import torch
import matplotlib.pyplot as plt

#生成测试数据
X = torch.rand(100, 1) * 10  # 生成一个100行一列的数据；该数据服从[0,10]的uniform分布
y = 3 * X + 10 + torch.randn(100, 1) * 3 # 计算其对应的y值；y也是100行1列的

class LinearRegression(torch.nn.Module):
    """
    模型需要继承 `torch.nn.Module`，在Pytorch中，模型都需要继承该类
    """
    def __init__(self):
        super().__init__()  # 初始化Module类
        """
        定义我们神经网络的第一层（线性层）。其接受的重要三个参数：
        in_features: 输入神经元的个数
        out_features：输出神经元的个数
        bias：是否包含偏置
        """
        self.linear = torch.nn.Linear(in_features=1, out_features=1, bias=True)

    def forward(self, x):
        """
        前向传播计算神经网络的输出
        """
        predict = self.linear(x)
        return predict

model = LinearRegression() # 初始化模型

#构建特征
"""
torch.optim.SGD 接受几个重要的参数：
- params: 模型参数
- lr: 学习率
"""
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# 这里可以看下模型参数
for param in model.parameters(): # 因为模型有多个参数，所以model.parameters会返回一个可迭代的对象
    print(param)

#定义损失函数MSE
loss_function = torch.nn.MSELoss()

#训练模型
for epoch in range(10000): # 训练10000次

    predict_y = model(X) #将X带入模型，其会自动调用前向传递，计算出每个x对应的y值

    loss = loss_function(predict_y, y) #通过损失函数计算损失

    loss.backward() #进行反向传播

    optimizer.step() #更新权重

    optimizer.zero_grad() #清空optimizer的梯度

#参看拟合情况
plt.scatter(X, y)
plt.plot(X, model(X).detach().numpy(), color='red')
plt.show()