pytorch 线性回归模型代码

以下是使用PyTorch实现线性回归模型的代码：

import torch

# 准备数据
X = torch.tensor([[1], [2], [3], [4]])  # 输入特征
y = torch.tensor([[2], [4], [6], [8]])  # 目标值

# 将我们的模型构造为一个类
class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)  # 构造一个对象，包含权重和偏置

    def forward(self, x):
        y_pred = self.linear(x)  # 在这里计算 w * x + b 线性模型
        return y_pred
model = LinearModel()

# 定义损失函数和优化器
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(100):
    # 前馈计算
    y_pred = model(X)
    
    # 计算损失
    loss = criterion(y_pred, y)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# 打印训练后的模型参数
print("训练后的模型参数：")
print("权重：", model.linear.weight.item())
print("偏置：", model.linear.bias.item())