pytorch 实现简单的线性回归模型

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)