深度学习之用PyTorch实现线性回归-Toy模板网

这篇具有很好参考价值的文章主要介绍了深度学习之用PyTorch实现线性回归。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

代码

# 调用库
import torch

# 数据准备
x_data = torch.Tensor([[1.0], [2.0], [3.0]])  # 训练集输入值
y_data = torch.Tensor([[2.0], [4.0], [6.0]])  # 训练集输出值


# 定义线性回归模型
class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()  # 调用父类构造函数
        self.linear = torch.nn.Linear(1, 1)  # 实例化torch库nn模块的Linear类

    def forward(self, x):
        """
        前馈运算
        :param x: 输入值
        :return: 线性回归预测结果
        """
        y_pred = self.linear(x)
        return y_pred


model = LinearModel()  # 实例化LinearModel类

criterion = torch.nn.MSELoss(size_average=False)  # 损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化器——梯度下降SGD
# optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # 优化器——Adam
# optimizer = torch.optim.Adamax(model.parameters(), lr=0.01)  # 优化器——Adamax


# 训练过程
for epoch in range(1000):  # epoch：训练轮次
    y_pred = model(x_data)
    loss = criterion(y_pred, y_data)
    print(epoch, 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)