import torch
import matplotlib.pyplot as plt

x_data = torch.tensor([[1.0], [2.0], [3.0]])  # x_data是一个张量
y_data = torch.tensor([[2.0], [4.0], [6.0]])  # y_data是一个张量


# 定义一个线性回归模型
class LinearModel(torch.nn.Module):  # 继承torch.nn.Module
    def __init__(self):
        super(LinearModel, self).__init__()  # 等价于nn.Module.__init__(self)
        self.linear = torch.nn.Linear(1, 1)  # 输入和输出的维度都是1

    def forward(self, x): # forward函数是必须要有的，用来构建计算图
        y_pred = self.linear(x)  # forward
        return y_pred


model = LinearModel()  # 实例化一个模型
criterion = torch.nn.MSELoss(size_average=False)  # 损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化器 lr为学习率
# 记录loss
loss_list = []
for epoch in range(100): # 训练100次
    y_pred = model(x_data) # forward
    loss = criterion(y_pred, y_data) # compute loss
    print(epoch, loss.item()) # 打印loss
    optimizer.zero_grad()  # 梯度清零
    loss.backward() # backward
    optimizer.step() # update
    # 记录loss
    loss_list.append(loss.item())

    print("predict (after training)", 4, model.forward(torch.tensor([4.0])).item()) # 预测

#绘制损失和训练次数的关系

plt.plot(range(100),loss_list)
plt.xlabel("epoch")
plt.ylabel("loss")
plt.show()