import torch

x_data = torch.tensor([[1.0], [2.0], [3.0]])
y_data = torch.tensor([[2.0], [4.0], [6.0]]) # 二维张量：每行是一个样本，每列是一个特征。表示3个样本，每个样本1个特征。

class Linearmodel(torch.nn.Module):
    def __init__(self):
        super().__init__() # （新版）
        # super(Linearmodel, self).__init__() # 必须调用父类初始化（旧版）
        self.linear = torch.nn.Linear(1, 1) # 创建线性层：输入1维，输出1维

    def forward(self, x):
        y_pred = self.linear(x) # 前向传播：y = wx + b
        return y_pred

model = Linearmodel() # 实例化模型

# criterion = torch.nn.MSELoss(size_average=False) # 损失函数（旧版）
criterion = torch.nn.MSELoss(reduction='sum') # 损失函数（总和，非平均）
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 优化器

for epoch in range(100):
    y_pred = model(x_data) # 前向传播
    loss = criterion(y_pred, y_data) # 计算损失
    print("epoch:", epoch, "loss:", loss.item())
    optimizer.zero_grad() # ☆梯度清零
    loss.backward() # 反向传播（自动计算梯度）
    optimizer.step() # 参数更新（自动使用梯度更新所有参数）
print('w = ', model.linear.weight.item()) # 获取训练后的权重
print('b = ', model.linear.bias.item()) # 获取训练后的偏置

x_test = torch.tensor([[4.0]]) # 测试输入，注意保持二维形状
y_test = model(x_test) # 前向传播得到预测
print("y_pred:", y_test.data) # 打印预测值