pytorch 实现简单的线性回归模型_pytorch nn.module 简单的线性回归-优快云博客

本文链接：https://blog.youkuaiyun.com/VanChouL/article/details/122144788

该博客展示了如何使用PyTorch构建一个简单的线性模型，包括模型定义、前向传播、损失计算、梯度清零、反向传播和优化器更新参数的过程。通过训练样本（x_data, y_data）进行迭代训练，并在训练结束后输出模型的权重和偏置，以及对测试样本的预测结果。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch


# 流程：1.设计模型
#      2.算y_hat
#      3.算loss
#      4.梯度清零
#      5.反向传播求梯度（变量对loss的导数）
#      6.优化器（更新参数）
#      7.返回第二步
class LinearModel(torch.nn.Module):
    def __init__(self):
        # 调用父类的方法
        super().__init__()

        # 两个参数分别是输入维度，和输出维度
        # nn 表示神经元网络，Liner返回一个对象
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        # 对象后加括号，表示可调用对象callable对象
        y_hat = self.linear(x)
        return y_hat


# 父类Module是可调用类，可以直接model（x）调用foward方法
model = LinearModel()

# 求loss，不求均值
criterion = torch.nn.MSELoss(size_average=False)

# 优化器，更新w和b，parameters是调用父类的方法
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练样本
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])

for epoch in range(100):
    y_hat = model(x_data)
    loss = criterion(y_hat, y_data)
    print(epoch, loss)

    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_hat = ", y_test.data)