05用PyTorch实现线性回归

用Pytorch实现线性回归

import torch
import numpy as np

x_data = torch.randint(1, 6, (3, 1)).float()
y_data = (x_data * 2 + 3.0).float()

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)
        return y_pred

model = LinearModel()  # 实例化

# 构建损失函数
criterion = torch.nn.MSELoss(size_average=True, reduction='mean')

# 构建优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)   # model.parameters自动计算parameter

for epoch in range(1000):
    y_pred = model(x_data)
    loss = criterion(y_pred, y_data)
    print(f"epoch: {epoch}\tloss: {loss.item()}")  # loss是对象，print的时候会自动调用__str__函数

    optimizer.zero_grad()   # 梯度清零
    loss.backward()  # 反向计算梯度
    optimizer.step()   # 梯度更新


print("w = ", model.linear.weight.item())
print("b = ", model.linear.bias.item())


# 预测
x_test = torch.tensor([[2.0]])
print(model(x_test))